diff --git a/README.md b/README.md index f652c70..f6d5ef8 100644 --- a/README.md +++ b/README.md @@ -1,32 +1,44 @@ # tom-alertstreams -`tom-alertstreams` is a reusable TOM Toolkit app for listening to kafka streams. - -`tom-alertstreams` provides a management command, `readstreams`. There are no `urlpatterns`, -no Views, and no templates. The `readstreams` management command reads the `settings.py` `ALERT_STREAMS` -configuration and starts listening to each configured Kafka stream. It is not expected -to return, and is intended to run along side your TOM's server component. The `ALERT_STREAMS` -configuration (see below) tells `readstreams` what streams to access, how to access them, -what topics to listen to, and what to do with messages that arrive on a given topic. +`tom-alertstreams` is a reusable [TOM Toolkit](https://tom-toolkit.readthedocs.io/) app +that listens to Kafka-based astronomical alert streams, persists incoming alerts to a +database, and displays them on a "Recent Alerts" page. + +## Features + +- **`readstreams` management command** — connects to configured Kafka streams, dispatches + each incoming alert to a topic-specific handler, and runs indefinitely alongside your + TOM's web server. +- **`Alert` model** — a FIFO-queued Django model that stores normalized alert fields + (stream name, topic, timestamp, RA/Dec, magnitude, raw payload). Older rows are + automatically pruned when the per-stream limit is exceeded. +- **Recent Alerts page** — an HTMX-driven filterable table at `/alertstreams/recent/` + that auto-registers via TOM Toolkit's AppConfig integration points (no manual URL + include or navbar edit required). +- **`save_alert_to_database` handler** — a ready-made `TOPIC_HANDLERS` value that + normalizes and persists any incoming alert to the database. + +--- ## Installation -1. Install the package into your TOM environment, specifying which alert streams you want to use: +1. Install the package into your TOM environment, specifying which alert streams you want: + ```bash # Install support for specific streams (recommended) - pip install tom-alertstreams[gcn] # For GCN Classic over Kafka - pip install tom-alertstreams[hopskotch] # For SCiMMA Hopskotch - pip install tom-alertstreams[antares] # For ANTARES - pip install tom-alertstreams[fink] # For Fink - - # Install multiple streams + pip install tom-alertstreams[gcn] # GCN Classic over Kafka + pip install tom-alertstreams[hopskotch] # SCiMMA Hopskotch + pip install tom-alertstreams[antares] # ANTARES + pip install tom-alertstreams[fink] # Fink + + # Multiple streams at once pip install tom-alertstreams[gcn,hopskotch,antares] - - # Or install all supported streams + + # All supported streams pip install tom-alertstreams[all-streams] - ``` + ``` -2. In your project `settings.py`, add `tom_alertstreams` to your `INSTALLED_APPS` setting: +2. Add `tom_alertstreams` to `INSTALLED_APPS` in your `settings.py`: ```python INSTALLED_APPS = [ @@ -35,184 +47,223 @@ what topics to listen to, and what to do with messages that arrive on a given to ] ``` -At this point you can verify the installation by running `./manage.py` to list the available -management commands and see +3. Run migrations to create the `Alert` table: + + ```bash + python manage.py migrate + ``` + +That's it. If your project uses the TOM Toolkit base URLs (`include('tom_common.urls')`), +the Recent Alerts page is now available at `/alertstreams/recent/` and the "Recent Alerts" +navbar link appears automatically — no additional URL or template changes required. - ```bash - [tom_alertstreams] - readstreams - ``` -in the output. +Verify the installation by checking that `readstreams` appears in `./manage.py` output: + +```bash +[tom_alertstreams] + readstreams +``` + +--- ## Configuration -Each Kafka stream that your TOM listens to (via `readstreams`) will have a configuration dictionary -in your `settings.py` `ALERT_STREAMS`. `ALERT_STREAMS` is a list of configuration dictionaries, one -dictionary for each Kafka stream. Here's an example `ALERT_STREAMS` configuration for three Kafka streams: -[SCiMMA Hopskotch](https://scimma.org/hopskotch.html), -[GCN Classic over Kafka](https://gcn.nasa.gov/quickstart), and -[ANTARES](https://nsf-noirlab.gitlab.io/csdc/antares/client/). +Add an `ALERT_STREAMS` list to your `settings.py`. Each entry is a dict with three keys: + +| Key | Type | Description | +|-----|------|-------------| +| `ACTIVE` | `bool` | Set `False` to disable a stream without removing its config. | +| `NAME` | `str` | Dotted path to an `AlertStream` subclass. | +| `OPTIONS` | `dict` | Stream-specific connection and topic-handler configuration. | + +The `OPTIONS` dict is validated by a Pydantic model on startup — missing required fields +raise descriptive errors before `readstreams` attempts any network connection. + +### Optional settings ```python -ALERT_STREAMS = [ - { - 'ACTIVE': True, - 'NAME': 'tom_alertstreams.alertstreams.hopskotch.HopskotchAlertStream', - 'OPTIONS': { - 'URL': 'kafka://kafka.scimma.org/', - # The hop-client requires that the GROUP_ID prefix match the SCIMMA_AUTH_USERNAME - 'GROUP_ID': os.getenv('SCIMMA_AUTH_USERNAME', "") + '-' + 'uniqueidforyourapp12345', - 'USERNAME': os.getenv('SCIMMA_AUTH_USERNAME', None), - 'PASSWORD': os.getenv('SCIMMA_AUTH_PASSWORD', None), - 'START_POSITION': 'LATEST', # Optional: EARLIEST or LATEST (defaults to LATEST) - 'TOPIC_HANDLERS': { - 'sys.heartbeat': 'tom_alertstreams.alertstreams.hopskotch.heartbeat_handler', - 'tomtoolkit.test': 'tom_alertstreams.alertstreams.hopskotch.alert_logger', - 'hermes.test': 'tom_alertstreams.alertstreams.hopskotch.alert_logger', - 'hermes.*': 'regex match public topics here, requires * handler to be defined' - '*': 'default_handler_here' - }, +# Enable the Recent Alerts web page and navbar link (default: False). +# When False (or absent), the /alertstreams/recent/ URL is not registered and +# the navbar link does not appear. The Alert model, migrations, and +# save_alert_to_database handler remain available regardless of this setting — +# only the web display layer is gated. Set to True if your TOM should display +# a live table of incoming alerts. +SHOW_RECENT_ALERTS = True + +# Maximum number of alerts to retain per stream (default: 100). +# Only relevant when SHOW_RECENT_ALERTS = True and save_alert_to_database is used. +ALERTSTREAMS_RECENT_COUNT = 100 +``` + +--- + +## Stream configuration reference + +### SCiMMA Hopskotch + +```python +{ + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.hopskotch.HopskotchAlertStream', + 'OPTIONS': { + 'URL': 'kafka://kafka.scimma.org/', + # GROUP_ID must be prefixed with your SCiMMA username. + 'GROUP_ID': os.environ.get('SCIMMA_AUTH_USERNAME', '') + '-my-tom', + 'USERNAME': os.environ.get('SCIMMA_AUTH_USERNAME', ''), + 'PASSWORD': os.environ.get('SCIMMA_AUTH_PASSWORD', ''), + 'START_POSITION': 'LATEST', # optional: 'LATEST' (default) or 'EARLIEST' + 'TOPIC_HANDLERS': { + 'sys.heartbeat': 'tom_alertstreams.alertstreams.hopskotch.heartbeat_handler', + # Wildcard patterns: 'hermes.*' matches all hermes.* topics + 'hermes.*': 'tom_alertstreams.alertstreams.handlers.save_alert_to_database', + # '*' matches ALL public topics not covered by a more specific entry + '*': 'tom_alertstreams.alertstreams.hopskotch.alert_logger', }, }, - { - 'ACTIVE': True, - 'NAME': 'tom_alertstreams.alertstreams.gcn.GCNClassicAlertStream', - # The keys of the OPTIONS dictionary become (lower-case) properties of the AlertStream instance. - 'OPTIONS': { - # see https://github.com/nasa-gcn/gcn-kafka-python#to-use for configuration details. - 'GCN_CLASSIC_CLIENT_ID': os.getenv('GCN_CLASSIC_CLIENT_ID', None), - 'GCN_CLASSIC_CLIENT_SECRET': os.getenv('GCN_CLASSIC_CLIENT_SECRET', None), - 'DOMAIN': 'gcn.nasa.gov', # optional, defaults to 'gcn.nasa.gov' - 'CONFIG': { # optional - # 'group.id': 'tom_alertstreams-my-custom-group-id', - # 'auto.offset.reset': 'earliest', - # 'enable.auto.commit': False - }, - 'TOPIC_HANDLERS': { - 'gcn.classic.text.LVC_INITIAL': 'tom_alertstreams.alertstreams.alertstream.alert_logger', - 'gcn.classic.text.LVC_PRELIMINARY': 'tom_alertstreams.alertstreams.alertstream.alert_logger', - 'gcn.classic.text.LVC_RETRACTION': 'tom_alertstreams.alertstreams.alertstream.alert_logger', - }, +}, +``` + +Credentials: [hop.scimma.org](https://hop.scimma.org/) + +### GCN Classic over Kafka + +```python +{ + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.gcn.GCNClassicAlertStream', + 'OPTIONS': { + 'GCN_CLASSIC_CLIENT_ID': os.environ.get('GCN_CLASSIC_CLIENT_ID', ''), + 'GCN_CLASSIC_CLIENT_SECRET': os.environ.get('GCN_CLASSIC_CLIENT_SECRET', ''), + 'DOMAIN': 'gcn.nasa.gov', # optional, default shown + 'KAFKA_CONFIG': {}, # optional dict passed to the Confluent Kafka Consumer + 'TOPIC_HANDLERS': { + 'gcn.classic.text.LVC_INITIAL': 'tom_alertstreams.alertstreams.handlers.save_alert_to_database', + 'gcn.classic.text.LVC_PRELIMINARY': 'tom_alertstreams.alertstreams.gcn.alert_logger', }, }, - { - 'ACTIVE': True, - 'NAME': 'tom_alertstreams.alertstreams.antares.AntaresAlertStream', - 'OPTIONS': { - 'API_KEY': os.getenv('ANTARES_API_KEY'), - 'API_SECRET': os.getenv('ANTARES_API_SECRET'), - 'GROUP': os.getenv('ANTARES_GROUP_ID'), - 'TOPIC_HANDLERS': { - 'extragalactic_staging': 'tom_antares.alertstream_handlers.handle_alert', - } +}, +``` + +Credentials: [gcn.nasa.gov/quickstart](https://gcn.nasa.gov/quickstart) + +### ANTARES + +```python +{ + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.antares.AntaresAlertStream', + 'OPTIONS': { + 'API_KEY': os.environ.get('ANTARES_API_KEY', ''), + 'API_SECRET': os.environ.get('ANTARES_API_SECRET', ''), + 'GROUP': 'my-tom-consumer-group', # optional, default: 'tom-alertstreams' + 'SSL_CA_LOCATION': None, # optional path to CA cert, default: None + 'ENABLE_AUTO_COMMIT': True, # optional, default: True + 'TOPIC_HANDLERS': { + 'nasa-ztf-test': 'tom_alertstreams.alertstreams.handlers.save_alert_to_database', }, - } -] + }, +}, ``` -The configuration dictionary for each `AlertStream` subclass will contain these key-value pairs: -* `ACTIVE`: Boolean which tells `readstreams` to access this stream. Should be `True`, unless you want to -keep a configuration dictionary, but ignore the stream. -* `NAME`: The name of the `AlertStream` subclass that implements the interface to this Kafka stream. `tom_alertstreams` -will provide `AlertStream` subclasses for major astromical Kafka streams. See below for instructions on Subclassing -the `AlertStream` base class. -* `OPTIONS`: A dictionary of key-value pairs specific to the`AlertStream` subclass given by `NAME`. The doc string for -the `AlertStream` subclass should document what is expected. Typically, a URL, authentication information, and a -dictionary, `TOPIC_HANDLERS`, will be required. See "Subclassing `AlertStream`" below. The `AlertStream` subclass will -convert the key-value pairs of the `OPTIONS` dictionary into properties (and values) of the `AlertStream` subclass -instance. - * The hopskotch alert stream supports a wildcard of `*` for an alert handler topic name. If specified, ALL public topics will be subscribed and use that handler function. A directly specified topic handler will always be used before the `*` handler for any topic that is covered twice. - -### Getting Kafka Stream Credentials -As part of your `OPTIONS` for each Kafka stream, you need to configure access credentials. Visit these links -to get credentials for [Hopskotch](https://hop.scimma.org/) and [GCN Classic over Kafka](https://gcn.nasa.gov/quickstart). -Set the environment variables with the username and passwords obtained. Do not check them in to your code repository. - - -## Alert Handling - -Assuming that an `AlertStream` subclass exists for the Kafka stream of interest, -the keys of the `TOPIC_HANDLERS` dictionary are the topics that will be subscribed to. The values -of the `TOPIC_HANDLERS` dictionary specify alert handling methods that will be imported and called -for each alert recieved on that topic. An example is provided, -`tom_alerts.alertstreams.alertstream.alert_logger`, which simply logs the alert. - -To customize this behaviour according to the needs of your TOM, define an alert handling function for each -topic that you wish to subscribe to. Your `TOPIC_HANDLERS` dictionary will have a an entry for each topic -whose key is the topic name and whose value is a string indicating the dot-path to the alert handling function. -When the `AlertStream` subclass is instanciated, the `OPTIONS` dictionary is read and an `alert_handler` -dictionary is created. It is keyed by topic name and it's values are the imported callable functions specified by the -dot-path strings. `readstreams` will call the alert handler for each alert that comes in on the topic. The signiture -of the alert handling function is specific to the `AlertStream` subclasss. - -## Subclassing `AlertStream` - -Ideally, As a TOM developer, there is already an `AlertStream`-subclass for the alert stream that you -want your TOM to listen to. If so, you need only to configure your TOM to use it in `settings.py` -`ALERT_STREAMS`. If you must implement your own `AlertStream` subclass, please get in touch. In the meantime, here's a brief outline: - -1. Create subclass of `AlertStream`. - -2. Create `required_keys` and `allowed_keys` class variables in your `AlertStream`-subclass. - - These are lists of strings refering to the keys of the `OPTIONS` dictionary. The purpose of these is to - help TOM developers using your `AlertStream`-subclass with the key-value pairs in their `ALERT_STREAMS` - `OPTIONS` configuration dictionary. - -3. Implement the `listen()` method. - - This method will be called by the `readstreams` management command and is not expected to return. It - should instanciate your consumer, subscribe to the topics configured in `ALERT_STREAMS`, and start - consuming. The detail of this will depend on the kafka-client used. See `alertstreams.gcn.listen()` - and `alertstreams.hopskotch.listen()` for examples to follow. - - The loop which consumes messages in your `listen()` method should extract the topic from each message - and call `self.alert_handler[topic]()` with the message or message-derived arguments specific to your - kafka client. Users of your `AlertStream`-subclass will write these topic-specific alert handling methods - and configure them in the `TOPIC_HANLDERS` dictionary of their `ALERT_STREAMS` configuration. - The `AlertStream` base class will set up the `alert_handler` dictionary according to your users' - configuration. It helps your users to provide an example `alert_hander()` function in your module as - an example. (Again, see `alertstreams.gcn.listen()` and `alertstreams.hopskotch.listen()`, their - configurations in `settings.py`, and the `alertstreams.gcn.alert_logger()` and - `alertstreams.hopskotch.alert_logger() methods, for example). +Credentials: [antares.noirlab.edu](https://antares.noirlab.edu) + +### Fink +Install: `pip install tom-alertstreams[fink]` + +Fink runs a separate Kafka broker and web portal per survey, and its topics follow +the `_` convention (e.g. `fink_sn_candidates_ztf`, +`fink_sn_candidates_lsst`). `FINK_SERVER`, `FINK_SURVEY`, and the topic suffixes must +all agree — a mismatch (e.g. `FINK_SURVEY='lsst'` with `_ztf` topics) is rejected at +startup, because it would route alerts to the wrong parser and silently drop them. Read +one survey per stream; configure two `ALERT_STREAMS` entries to read both. + +```python +{ + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.fink.FinkAlertStream', + 'OPTIONS': { + 'FINK_USERNAME': os.environ.get('FINK_USERNAME', ''), + 'FINK_GROUP_ID': os.environ.get('FINK_GROUPID', ''), + 'FINK_SERVER': 'kafka-ztf.fink-broker.org:24499', # LSST: kafka-lsst.fink-broker.org:24499 + 'FINK_SURVEY': 'ztf', # 'ztf' or 'lsst' + 'TOPIC_HANDLERS': { + 'fink_sn_candidates_ztf': 'tom_alertstreams.alertstreams.alertstream.save_alert_to_database', + }, + }, +}, +``` + +Credentials: register at [fink-broker.org](https://fink-broker.org). + +### ALeRCE, AMPEL, Babamul, Lasair, Pitt-Google + +These LSST broker stubs generate obviously-fake mock alerts (ra=0, dec=0, +magnitude=99, object IDs prefixed with `MOCK-`) for demonstration and development. +Replace with real implementations when the broker clients become available. + +```python +{ + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.alerce.AlerceAlertStream', + 'OPTIONS': { + 'TOPIC_HANDLERS': { + 'alerce-topic': 'tom_alertstreams.alertstreams.handlers.save_alert_to_database', + }, + }, +}, +# Similarly for ampel, babamul, lasair, pittgoogle — same OPTIONS structure. +``` + +--- + +## Running the alert listener + +Start `readstreams` alongside your Django development server: + +```bash +python manage.py readstreams +``` + +Each configured active stream runs in its own thread. `readstreams` is not expected to +return. In production, run it as a separate process managed by systemd, supervisor, or +your preferred process manager. + +--- + +## Alert persistence + +`save_alert_to_database` is a ready-made handler that normalizes an alert and saves it +to the `Alert` model. Use it as a `TOPIC_HANDLERS` value: + +```python +'TOPIC_HANDLERS': { + 'my.topic': 'tom_alertstreams.alertstreams.handlers.save_alert_to_database', +} +``` + +Rows older than `ALERTSTREAMS_RECENT_COUNT` (default 100) per stream are automatically +pruned after each save. The Recent Alerts page at `/alertstreams/recent/` displays +persisted alerts with HTMX-powered live filtering by stream name or full-text search. + +--- ## Local Development -### Getting Started - -1. Clone the repository: - ```bash - git clone https://github.com/TOMToolkit/tom_alertstreams.git - cd tom_alertstreams - ``` - -2. Create and activate your virtual environment: - ```bash - /path/to/your/python -m venv .venv - source .venv/bin/activate - ``` - -3. Install dependencies (choose one): - ```bash - # Using pip - install all streams for development - pip install -e .[all-streams] - # Or install specific streams only - pip install -e .[gcn,hopskotch] - - # Using poetry - install all streams - poetry install --all-extras - # Or install specific streams only - poetry install --extras gcn --extras hopskotch - - # Using uv - install all streams - uv pip install -e .[all-streams] - # Or install specific streams only - uv pip install -e .[gcn,hopskotch] - ``` - -4. Run the tests: - ```bash - python tom_alertstreams/tests/run_tests.py - ``` +```bash +git clone https://github.com/TOMToolkit/tom_alertstreams.git +cd tom_alertstreams + +# Create and activate a virtual environment (Python 3.9+) +python -m venv .venv +source .venv/bin/activate + +# Install with stream extras for development +pip install -e ".[all-streams]" + +# Run tests (requires a full TOM Toolkit environment) +python manage.py test tom_alertstreams --exclude-tag=canary +``` +For extending `tom_alertstreams` — writing custom handlers, subclassing `AlertStream`, +or overriding the Alert model — see the extension guide in the TOM Toolkit documentation. diff --git a/pyproject.toml b/pyproject.toml index 98e9cd7..eed5cea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,20 +36,30 @@ dynamic = ["version"] requires-python = ">=3.9.0,<=3.13" dependencies = [ "django <6", - "django-extensions >=3.2.3,<4" + "django-extensions >=3.2.3,<4", + "pydantic >=2.0,<3" ] [project.optional-dependencies] gcn = ["gcn-kafka >=0.3,<1.0"] hopskotch = ["hop-client >=0.10,<1.0"] antares = ["antares-client"] -fink = ["fink-client >=8.8,<9"] +babamul = ["babamul >=0.1,<1"] +fink = ["fink-client >=11.0,<12"] +ampel-ztf = ["ampel-ztf"] +ampel-lsst = ["ampel-lsst"] +ampel = ["ampel-ztf", "ampel-lsst"] +pittgoogle = ["pittgoogle-client >=0.3.22,<1.0"] all-streams = [ "gcn-kafka >=0.3,<1.0", "hop-client >=0.10,<1.0", "antares-client", - "fink-client >=8.8,<9", + "babamul >=0.1,<1", + "fink-client >=11.0,<12", + "ampel-ztf", + "ampel-lsst", + "pittgoogle-client >=0.3.22,<1.0" ] test = [ diff --git a/tom_alertstreams/admin.py b/tom_alertstreams/admin.py index 8c38f3f..78a7564 100644 --- a/tom_alertstreams/admin.py +++ b/tom_alertstreams/admin.py @@ -1,3 +1,10 @@ from django.contrib import admin -# Register your models here. +from tom_alertstreams.models import Alert + + +@admin.register(Alert) +class AlertAdmin(admin.ModelAdmin): + list_display = ('stream_name', 'alert_id', 'published_time', 'observation_time', 'object_id', 'magnitude', 'flux') + list_filter = ('stream_name',) + search_fields = ('alert_id', 'object_id') diff --git a/tom_alertstreams/alertstreams/alerce.py b/tom_alertstreams/alertstreams/alerce.py new file mode 100644 index 0000000..b92c52b --- /dev/null +++ b/tom_alertstreams/alertstreams/alerce.py @@ -0,0 +1,268 @@ +from __future__ import annotations + +import io +import logging +import random +import time +from datetime import datetime, timezone +from typing import Any, ClassVar + +from confluent_kafka import Consumer +from fastavro import reader as avro_reader + +from tom_alertstreams.alertstreams.alertstream import ( + AlertStream, AlertStreamConfig, NormalizedAlert, _mjd_to_datetime, +) + +logger = logging.getLogger(__name__) + +# Mock data constants — chosen to be unmistakably non-astronomical: +# (0, 0) is not a real survey pointing; 99.0 is the astronomical sentinel for "no data". +_MOCK_RA = 0.0 +_MOCK_DEC = 0.0 +_MOCK_MAGNITUDE = 99.0 +INTER_ALERT_SLEEP_MIN = 3600 # sixty minutes +INTER_ALERT_SLEEP_MAX = 3600 + + +# --------------------------------------------------------------------------- +# Pydantic configuration models +# --------------------------------------------------------------------------- + +class AlerceMockConfig(AlertStreamConfig): + """Pydantic configuration model for AlerceMockAlertStream. + + Inherits TOPIC_HANDLERS from AlertStreamConfig (a Pydantic BaseModel). + No additional fields needed for mock data generation. + """ + pass + + +class AlerceConfig(AlertStreamConfig): + """Pydantic configuration for AlerceAlertStream. + + ALeRCE applies ML classifiers to the survey alert stream (ZTF now, LSST in + preparation) and re-publishes the results over Kafka. Access requires credentials, + obtained by emailing alerce.broker@gmail.com — see + https://github.com/alercebroker/Kafka-Connection-Docs. + + Fields: + ALERCE_KAFKA_SERVER: Kafka bootstrap server. Defaults to the public ALeRCE broker. + ALERCE_GROUP_ID: Kafka consumer group ID (required). + ALERCE_USERNAME: SASL/SCRAM username issued by ALeRCE (required). + ALERCE_PASSWORD: SASL/SCRAM password issued by ALeRCE (required). + TOPIC_PREFIX: ALeRCE publishes a new topic per UTC day named '{prefix}_YYYYMMDD' + (e.g. 'lc_classifier_20260605'), each retained ~48 hours. We subscribe by + regex on this prefix so the consumer follows the rolling daily topics without + a config change. Use 'lc_classifier' (light-curve classifier) or + 'stamp_classifier' (stamp classifier). + TOPIC_HANDLERS: Inherited from AlertStreamConfig. Because the literal topic name + changes daily, the single entry here is keyed on TOPIC_PREFIX rather than a + concrete topic, e.g. + {'lc_classifier': 'tom_alertstreams.alertstreams.alertstream.save_alert_to_database'}. + """ + ALERCE_KAFKA_SERVER: str = 'kafka.alerce.science:9093' + ALERCE_GROUP_ID: str + ALERCE_USERNAME: str + ALERCE_PASSWORD: str + TOPIC_PREFIX: str = 'lc_classifier' + + +# --------------------------------------------------------------------------- +# Mock ALeRCE stream (for demo use without ALeRCE credentials) +# --------------------------------------------------------------------------- + +class AlerceMockAlertStream(AlertStream): + """Mock ALeRCE AlertStream that generates obviously-fake alerts. + + Demo fallback when ALeRCE Kafka credentials are not available. Uses the same mock + pattern as the other stub streams (sentinel coordinates, sentinel magnitude, 6-7 + minute sleep between alerts). STREAM_NAME='alerce' so it occupies the same dashboard + slot the real ALeRCE stream would. + """ + configuration_class = AlerceMockConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'alerce' + IS_MOCK: ClassVar[bool] = True + + def normalize_alert(self, raw_alert: dict, topic: str = '') -> NormalizedAlert: + """Map a mock ALeRCE alert dict to a NormalizedAlert. + + Args: + raw_alert: Dict produced by listen(); contains mock field values. + topic: Topic the alert was generated for. + + Returns: + NormalizedAlert populated from the mock dict fields. + """ + return NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic or raw_alert.get('topic', ''), + observation_time=None, # a mock alert has no real observation + published_time=datetime.fromisoformat(raw_alert['timestamp']), + alert_id=raw_alert['alert_id'], + object_id=raw_alert.get('object_id'), + ra=raw_alert.get('ra'), + dec=raw_alert.get('dec'), + magnitude=raw_alert.get('magnitude'), + raw_payload=raw_alert, + ) + + def listen(self) -> None: + """Generate mock ALeRCE alerts and dispatch to configured topic handlers. + + Loops indefinitely, emitting one mock alert per iteration with a random 6-7 + minute delay. Topics are round-robined if multiple are configured. + """ + counter = 0 + topics = list(self.config.TOPIC_HANDLERS.keys()) + + while True: + counter += 1 + topic = topics[counter % len(topics)] + timestamp = datetime.now(timezone.utc) + object_id = f'MOCK-{self.STREAM_NAME.upper()}-{counter:04d}' + alert_id = f'MOCK-{timestamp.strftime("%Y%m%d%H%M%S")}' + mock_alert: dict[str, Any] = { + 'alert_id': alert_id, + 'object_id': object_id, + 'topic': topic, + 'timestamp': timestamp.isoformat(), + 'ra': _MOCK_RA, + 'dec': _MOCK_DEC, + 'magnitude': _MOCK_MAGNITUDE, + 'mock': True, + } + logger.debug(f'AlerceMockAlertStream: mock alert {object_id}') + self.alert_handler[topic](mock_alert, alert_stream=self, topic=topic) + time.sleep(random.uniform(INTER_ALERT_SLEEP_MIN, INTER_ALERT_SLEEP_MAX)) + + +# --------------------------------------------------------------------------- +# Real ALeRCE stream +# --------------------------------------------------------------------------- + +class AlerceAlertStream(AlertStream): + """ALeRCE alert stream via Kafka. + + Consumes ALeRCE's classifier output with a raw confluent_kafka Consumer, following + the same lightweight pattern as LasairAlertStream. Authentication is SASL_PLAINTEXT + + SCRAM-SHA-256 with credentials issued by ALeRCE (email alerce.broker@gmail.com). + + Rolling daily topics: ALeRCE creates a new topic per UTC day, '{TOPIC_PREFIX}_YYYYMMDD', + each retained ~48 hours. We subscribe with a confluent_kafka regex ('^{prefix}_') so the + consumer follows the rolling topics automatically — no daily config change. Because the + literal topic name is therefore not known in advance, every message is dispatched to the + single handler configured under the TOPIC_PREFIX key (the single-handler approach + HopskotchAlertStream uses for its wildcard subscriptions). + + Messages are Avro; we decode them with fastavro (already installed as a fink-client + dependency). NOTE: this assumes container-framed Avro (schema embedded per message, as + ZTF-derived streams use). If ALeRCE sends schemaless/Confluent-wire Avro instead, the + decode in listen() must switch to fastavro.schemaless_reader with ALeRCE's published + schema — finalize against a real message once credentials are available. + + Configuration example (settings.py ALERT_STREAMS entry):: + + { + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.alerce.AlerceAlertStream', + 'OPTIONS': { + 'ALERCE_GROUP_ID': os.environ.get('ALERCE_GROUP_ID', ''), + 'ALERCE_USERNAME': os.environ.get('ALERCE_USERNAME', ''), + 'ALERCE_PASSWORD': os.environ.get('ALERCE_PASSWORD', ''), + 'TOPIC_PREFIX': 'lc_classifier', + 'TOPIC_HANDLERS': { + 'lc_classifier': 'tom_alertstreams.alertstreams.alertstream.save_alert_to_database', + }, + }, + } + """ + configuration_class = AlerceConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'alerce' + + def listen(self) -> None: + """Consume ALeRCE alerts from Kafka and dispatch to the configured handler. + + Subscribes by regex to the rolling daily topics and polls indefinitely. Avro + message values are decoded with fastavro; each decoded record is dispatched to the + single handler keyed on TOPIC_PREFIX. Errors propagate to AlertStream.run(), which + logs and restarts listen() — connection resilience lives there, not here. + """ + prefix = self.config.TOPIC_PREFIX + # The literal daily topic isn't in TOPIC_HANDLERS, so resolve the handler by the + # prefix key (falling back to the sole configured handler for robustness). + handler = self.alert_handler.get(prefix) or next(iter(self.alert_handler.values())) + + alerce_kafka_consumer = Consumer({ + 'bootstrap.servers': self.config.ALERCE_KAFKA_SERVER, + 'group.id': self.config.ALERCE_GROUP_ID, + 'security.protocol': 'SASL_PLAINTEXT', + 'sasl.mechanism': 'SCRAM-SHA-256', + 'sasl.username': self.config.ALERCE_USERNAME, + 'sasl.password': self.config.ALERCE_PASSWORD, + # Demo semantics: show only current alerts (the table is a latest-only FIFO). + 'auto.offset.reset': 'latest', + }) + # confluent_kafka treats a leading '^' as a subscription regex, matching every + # current and future '{prefix}_YYYYMMDD' topic. + alerce_kafka_consumer.subscribe([f'^{prefix}_']) + logger.info(f'{self.STREAM_NAME}: subscribed to ^{prefix}_ on {self.config.ALERCE_KAFKA_SERVER}') + + try: + while True: + msg = alerce_kafka_consumer.poll(timeout=20) + if msg is None: + continue # poll timeout with no message, retry + if msg.error(): + logger.warning(f'{self.STREAM_NAME} Kafka error: {msg.error()}') + continue + topic = msg.topic() + logger.info(f'{self.STREAM_NAME} received message on {topic}') + # ALeRCE messages are Avro; fastavro.reader handles container-framed Avro + # (schema embedded), yielding one record per alert. + for record in avro_reader(io.BytesIO(msg.value())): + handler(record, alert_stream=self, topic=topic) + finally: + alerce_kafka_consumer.close() + + def normalize_alert(self, raw_alert: dict, topic: str = '') -> NormalizedAlert: + """Extract common fields from an ALeRCE classifier record. + + ALeRCE classifier outputs are keyed by object id and candidate id, plus features + and class probabilities; coordinates and magnitude may or may not be present + depending on the classifier. We populate whatever is available. Field names follow + ALeRCE conventions: 'oid' (object id), 'candid' (candidate id), 'meanra'/'meandec', + 'lastmjd' (MJD). These should be confirmed against a real message. + + Args: + raw_alert: Avro-decoded ALeRCE record (a dict). + topic: The (daily) topic the alert arrived on. + + Returns: + NormalizedAlert with the available fields populated. + """ + record = dict(raw_alert) # fastavro yields plain dicts; copy for raw_payload safety + + object_id = record.get('oid') or record.get('aid') + candid = record.get('candid') + + # Timestamp from the object's most recent detection MJD, if present. + mjd = record.get('lastmjd') or record.get('firstmjd') + timestamp = _mjd_to_datetime(mjd) if mjd is not None else None + + ra = record.get('meanra', record.get('ra')) + dec = record.get('meandec', record.get('dec')) + magnitude = record.get('magpsf') + + return NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic, + observation_time=timestamp, + published_time=None, + alert_id=str(candid) if candid is not None else str(object_id or ''), + object_id=str(object_id) if object_id else None, + ra=float(ra) if ra is not None else None, + dec=float(dec) if dec is not None else None, + magnitude=float(magnitude) if magnitude is not None else None, + raw_payload=record, + ) diff --git a/tom_alertstreams/alertstreams/alertstream.py b/tom_alertstreams/alertstreams/alertstream.py index 11b812d..9797483 100644 --- a/tom_alertstreams/alertstreams/alertstream.py +++ b/tom_alertstreams/alertstreams/alertstream.py @@ -1,108 +1,489 @@ +from __future__ import annotations + import abc import logging +import time +from datetime import datetime, timezone +from typing import Any, Callable, ClassVar from django.conf import settings from django.core.exceptions import ImproperlyConfigured from django.utils.module_loading import import_string +from pydantic import BaseModel, ValidationError + +from tom_alertstreams.models import Alert + logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) -def get_default_alert_streams(): - """Return the AlertStreams configured in settings.py +# --------------------------------------------------------------------------- +# Julian Date / MJD helpers — shared across multiple alert stream modules +# --------------------------------------------------------------------------- + +def _jd_to_datetime(jd: float) -> datetime: + """Convert Julian Date to a timezone-aware UTC datetime. + + Uses the standard epoch offset: JD 2440587.5 = Unix epoch (1970-01-01 00:00:00 UTC). """ - try: - alert_streams = get_alert_streams(settings.ALERT_STREAMS) - except AttributeError as err: - raise ImproperlyConfigured(err) + unix_seconds = (jd - 2440587.5) * 86400.0 + return datetime.fromtimestamp(unix_seconds, tz=timezone.utc) - return alert_streams +def _mjd_to_datetime(mjd: float) -> datetime: + """Convert Modified Julian Date to a timezone-aware UTC datetime. -def get_alert_streams(alert_stream_configs: list): - """Return the AlertStreams configured in the given alert_stream_configs - (a list of configuration dictionaries ) + MJD = JD - 2400000.5, so we convert back to JD and delegate. + """ + return _jd_to_datetime(mjd + 2400000.5) + + +def is_in_hourly_window(timestamp_ms: int) -> bool: + """Return True if timestamp_ms falls in the first second of its UTC hour. + + A stateless once-per-hour throttle for high-frequency feeds. UTC hour boundaries + are exact multiples of 3_600_000 ms, so (timestamp_ms % 3_600_000) is the offset + into the current hour; < 1000 ms keeps just the one message in the hour's first + second. Used to thin firehose streams — GCN's ~1/sec heartbeat and Pitt-Google's + ~1/sec ztf-loop — down to a single saved alert per hour. - Use get_default_alert_streams() if you want the AlertStreams configured in settings.py. + Args: + timestamp_ms: A Unix-epoch timestamp in milliseconds (e.g. a Kafka message + timestamp or a Pub/Sub publishTime converted to ms). + + Returns: + True if the timestamp is within the first 1000 ms of its UTC hour. """ - alert_streams = [] # build and return this list of AlertStream subclass instances - for alert_stream_config in alert_stream_configs: - if not alert_stream_config.get('ACTIVE', True): - logger.debug(f'get_alert_streams - ignoring inactive stream: {alert_stream_config["NAME"]}') - continue # skip configs that are not active; default to ACTIVE - try: - klass = import_string(alert_stream_config['NAME']) - except ImportError: - msg = ( - f'The module (the value of the NAME key): {alert_stream_config["NAME"]} could not be imported. ' - f'Check your ALERT_STREAMS setting.' - ) - raise ImproperlyConfigured(msg) + return timestamp_ms % 3_600_000 < 1000 - alert_stream: AlertStream = klass(**alert_stream_config.get("OPTIONS", {})) - alert_streams.append(alert_stream) - return alert_streams +def is_in_minute_window(timestamp_ms: int) -> bool: + """Return True if timestamp_ms falls in the first second of its UTC minute. + + The per-minute analogue of is_in_hourly_window — a stateless once-per-minute throttle. + UTC minute boundaries are exact multiples of 60_000 ms, so (timestamp_ms % 60_000) is the + offset into the current minute; < 1000 ms keeps the one message in the minute's first + second. Thins a high-rate feed (Pitt-Google's ztf-loop / ztf-alerts) to ~1/min — a livelier + demo cadence than hourly, while still keeping acks fast enough to stay current. + + Args: + timestamp_ms: A Unix-epoch timestamp in milliseconds. + + Returns: + True if the timestamp is within the first 1000 ms of its UTC minute. + """ + return timestamp_ms % 60_000 < 1000 +# --------------------------------------------------------------------------- +# Typed alert intermediate +# --------------------------------------------------------------------------- + +class NormalizedAlert(BaseModel): + """Typed intermediate produced by AlertStream.normalize_alert(). + + Every AlertStream subclass's normalize_alert() method returns a NormalizedAlert. + Handlers that need to persist alerts (e.g. save_alert_to_database) rely on this + type so that one handler function works across all streams. + + Only stream_name and alert_id are required; every other field is optional because + not every stream provides the same metadata. The raw_payload preserves the full + original alert object (serialized to a dict) for handlers that need + stream-specific data not captured in the normalised fields. + + Three distinct clocks (all UTC), any of which may be absent for a given stream: + observation_time: when the telescope observed the source the alert is about + (e.g. a detection MJD/JD). Absent for non-observational alerts such as GCN + Circulars. + published_time: when the alert was issued/published by the broker or survey + (e.g. GCN Circular createdOn, Fink brokerEndProcessTimestamp). Distinct + from observation_time — the gap is the broker's processing latency. + (receipt time is the Alert.created column, set when we save the row.) + + Fields: + stream_name: Short canonical name of the stream (from AlertStream.STREAM_NAME). + alert_id: Stream-specific identifier for this alert. + topic: Kafka topic the alert arrived on. Empty string if not available. + observation_time: UTC observation datetime, if available. + published_time: UTC datetime the alert was issued by the broker/survey, if available. + object_id: Astronomical object identifier (e.g. ZTF object name), if available. + ra: Right ascension in decimal degrees, if available. + dec: Declination in decimal degrees, if available. + magnitude: Apparent magnitude, if available. ZTF streams populate this field. + flux: Flux in nanojansky, if available. LSST streams populate this field. + raw_payload: The full original alert as a plain dict for downstream use. + """ + stream_name: str + alert_id: str + topic: str = '' + observation_time: datetime | None = None + published_time: datetime | None = None + object_id: str | None = None + ra: float | None = None + dec: float | None = None + magnitude: float | None = None + flux: float | None = None + raw_payload: dict = {} + + +# --------------------------------------------------------------------------- +# Pydantic configuration models +# --------------------------------------------------------------------------- + +class AlertStreamConfig(BaseModel): + """Pydantic base configuration for all AlertStream subclasses. + + Inheriting from pydantic's BaseModel means that Pydantic validates required + fields, coerces types, and raises descriptive field-level ValidationError + messages when configuration is missing or incorrect — replacing the previous + manual required_keys / allowed_keys validation approach. + + Every stream must declare its topic-to-handler mapping. Subclass configs + inherit this and add stream-specific authentication and connection fields. + Handler values are dotted-path strings; they are imported at AlertStream + instantiation time by _process_topic_handlers(). + + Example subclass: + class MyStreamConfig(AlertStreamConfig): + USERNAME: str + PASSWORD: str + START_POSITION: str = 'LATEST' + """ + # Maps topic names to dotted-path strings of callable alert handler functions. + # Example: {'my.topic': 'myapp.handlers.save_alert_to_database'} + TOPIC_HANDLERS: dict[str, str] + + +# --------------------------------------------------------------------------- +# AlertStream abstract base class +# --------------------------------------------------------------------------- + class AlertStream(abc.ABC): - """Base class for specific alert streams like Hopskotch, GCNClassic, etc. - - * kwargs to __init__ is the OPTIONS dictionary defined in ALERT_STREAMS configuration - dictionary (for example, see settings.py). - * allowed_keys and required_keys should be defined as class properties in subclasses. - * The allowed_keys are turned into instance properties in __init__. - * Missing required_keys result in an ImproperlyConfigured Django exception. - - - To implmement subclass: - 1. define allowed_key, required_keys as class variables - - 2. implement listen() - this method probably doesn't return - 3. write your alert_handlers. which proably take and alert do something. - The HopskotchAlertStream.listen() method defines an 'alert_handlers' dictionary keyed by - alert topic with callable values (i.e call this method with alerts from this topic). - The GCNClassicAlertStream.listen() is another example. + """Abstract base class for Kafka alert stream implementations. + + To implement a new AlertStream subclass: + + 1. Define a Pydantic config model (subclass of AlertStreamConfig) that declares + all required and optional configuration fields. Pydantic handles validation + and produces error messages for misconfigured streams. + + 2. Set class variables: + configuration_class = MyStreamConfig + STREAM_NAME = 'mystream' # short canonical name, written to Alert.stream_name + + 3. Override normalize_alert(raw_alert, topic='') -> NormalizedAlert to extract + stream-specific fields (ra, dec, magnitude, object_id, etc.). This returns + a basic (Pydantic BaseModel subclass) NormaizedAlert that can be consistently + used reguardless of which stream the alert came from. + + 4. Implement listen() -> None. This method is not expected to return. It should: + a. Connect to the Kafka stream using credentials from self.config + b. Subscribe to the topics in self.config.TOPIC_HANDLERS + c. Dispatch each incoming alert to its handler. Alert handlers shoud + use the following function signiture: + self.alert_handler[topic](raw_alert, alert_stream=self, topic=topic) + Include any stream-specific extras as named keyword args (e.g. metadata=metadata). + Handlers use **kwargs to absorb extras they do not need. + + The alert_stream=self argument provides dependency injection, allowing the + same generic handler function (e.g. save_alert_to_database) to serve all + streams because it receives the stream instance and can call the AlertStream's + get_normalization_function() to obtain stream-specific parsing logic without + knowing which stream it is working with. """ + # alertstream.AlertStreamConfig is a Pydantic BaseModel subclass + # the settings.ALERT_STREAMS configuration dictionary will validated according + # to the AlertStreamConfig subclass specified here. + configuration_class: ClassVar[type[AlertStreamConfig]] - def __init__(self, *args, **kwargs) -> None: - super().__init__() + # Short canonical name written to Alert.stream_name. Must be unique across + # all configured streams. Used by the presenter registry in tables.py to + # look up the appropriate AlertStreamPresenter for URL construction. + STREAM_NAME: ClassVar[str] - # filter the kwargs by allowed keys and add them as properties to AlertStream instance - self.__dict__.update((k.lower(), v) for k, v in kwargs.items() if k in self.allowed_keys) + # True on mock/stub streams that generate simulated demo alerts (rather than + # connecting to the real broker). The dashboard uses this to indicate visually + # streams that aren't showing real data. + IS_MOCK: ClassVar[bool] = False - missing_keys = set(self.required_keys) - set(kwargs.keys()) - if missing_keys: - msg = ( - f'The following required keys are missing from the configuration OPTIONS of ' - f'{self._get_stream_classname()}: {list(missing_keys)} ; ' - f'These keys were found: {list(kwargs.keys())} ; ' - f'Check your ALERT_STREAMS setting.' - ) - raise ImproperlyConfigured(msg) + # Seconds run() waits before restarting listen() after it returns or raises. + # This is the reconnect backoff shared by every stream; a subclass may override + # it for a broker that needs a longer delay between connection attempts. + RESTART_DELAY_SECONDS: ClassVar[float] = 30.0 + + def __init__(self, **kwargs: Any) -> None: + # read and validate the alertstream configuration + self.config: AlertStreamConfig = self.configuration_class(**kwargs) - self.alert_handler = self._process_topic_handlers() + # Convert TOPIC_HANDLERS dotted-path strings to callable functions. + self.alert_handler: dict[str, Callable] = self._process_topic_handlers() def _get_stream_classname(self) -> str: + """Return the qualified class name of this AlertStream subclass. + + This is just a way to get the name of the subclass. + """ return type(self).__qualname__ - def _process_topic_handlers(self): - """ convert the TOPIC_HANDLERS values in to callable functions in - the returned message_handler dictionary. (keyed by topic, value is callable) + def _process_topic_handlers(self) -> dict[str, Callable]: + """Import and return handler callables from the TOPIC_HANDLERS configuration. + + This is a step in the AlertStream instanciation: + + In settings.py, the configuration dictionary TOPIC_HANDLER dictionary + for each stream maps a topic to a dotted-path string specifying the alert + handler for that topic's alerts. This method converts the dotted-path string + to a Callable. + + Returns: + A dict mapping topic name strings to callable handler functions. + + Raises: + ImproperlyConfigured: if any handler dotted-path cannot be imported. """ alert_handler = {} - for topic, callable_string in self.topic_handlers.items(): - # convert string from TOPIC_HANDLERS in to a callable function in - # the message_handler dictionary (both keyed by topic string) - alert_handler[topic] = import_string(callable_string) + for topic, callable_string in self.config.TOPIC_HANDLERS.items(): + try: + alert_handler[topic] = import_string(callable_string) + except ImportError as err: + msg = ( + f'Could not import handler "{callable_string}" for topic "{topic}" ' + f'in {self._get_stream_classname()}. Check your TOPIC_HANDLERS setting. ' + f'Error: {err}' + ) + raise ImproperlyConfigured(msg) return alert_handler @abc.abstractmethod - def listen(self): - """Listen at the steam and dispatch alerts to handlers. Subclass extentions of - this method are not expected to return. See hopskotch.py and gcn.py for example - implementations. + def normalize_alert(self, raw_alert: Any, topic: str = '') -> NormalizedAlert: + """Convert a raw stream-specific alert object to a NormalizedAlert. + + Every AlertStream subclass must implement this method. The returned + NormalizedAlert is a Pydantic BaseModel subclass. As such, it is typed + and validated (vs a dictionary of unvalidated key and values). + + While the stream is known by virtue of the AlertSteam subclass implementing + this method, the topic argument is provided by the listen() loop and should + be passed through to the NormalizedAlert.topic field. + + Args: + raw_alert: The stream-specific alert object received from listen(). + topic: The Kafka topic this alert arrived on. + + Returns: + A NormalizedAlert with as many fields populated as the stream supports. + """ + pass # implement me + + def get_normalization_function(self) -> Callable: + """Return the normalization callable for this stream. + + By default returns self.normalize_alert. Override this method to substitute + a completely different normalization implementation without subclassing, for + example to use a function defined outside this class/subclass (e.g. in your + `custom_code` or other INSTALLED_APP).) + + Returns: + A callable with signature (raw_alert, topic='') -> NormalizedAlert. + """ + return self.normalize_alert + + @abc.abstractmethod + def listen(self) -> None: + """Consume alerts for a single connect-and-consume session. + + In normal operation this method does not return — it connects once and + consumes indefinitely. It IS, however, allowed to raise on a broker drop, + auth failure, deserialization error, etc.: run() supervises listen() and + restarts it after a backoff, so implementations should NOT add their own + reconnect/restart loops. Connection resilience lives once, in run(). + + Implementations should: + 1. Connect to the Kafka stream using credentials from self.config + 2. Subscribe to the topics in self.config.TOPIC_HANDLERS (the topic + keys are also available via self.alert_handler.keys()) + 3. For each incoming alert, dispatch to the handler using the unified + calling convention: + + self.alert_handler[topic](raw_alert, alert_stream=self, topic=topic) + + The alert_stream=self argument injects this AlertStream instance so that + handlers can call get_normalization_function() without knowing the specific + stream type (dependency injection via keyword argument). + + Pass any additional stream-specific context as named keyword arguments: + + self.alert_handler[topic](raw_alert, alert_stream=self, topic=topic, + metadata=metadata) # Hopskotch example + + Handlers absorb extras they do not need via **kwargs. + """ + pass # implement me in your subclass + + def run(self) -> None: + """Supervise listen(), restarting it forever so a stream can't die silently. + + readstreams launches this (not listen() directly) in one thread per stream. + Because each listen() runs in a bare Thread with no restart, an exception that + escapes listen() would otherwise kill that one stream permanently — and + silently — while the other streams keep running. + + This wrapper logs any failure and restarts listen() after RESTART_DELAY_SECONDS, + giving every stream automatic reconnect-on-error for free. Only Exception is + caught, so KeyboardInterrupt / SystemExit still propagate for a clean shutdown. """ - pass + while True: + try: + self.listen() + # listen() is documented as not returning in normal operation; if it + # does, the session ended (e.g. the consumer was closed) — restart it. + logger.warning( + f'{self.STREAM_NAME}: listen() returned; restarting in {self.RESTART_DELAY_SECONDS}s.' + ) + except Exception as exc: + logger.exception( + f'{self.STREAM_NAME}: listen() failed ({exc.__class__.__name__}: {exc}); ' + f'restarting in {self.RESTART_DELAY_SECONDS}s.' + ) + time.sleep(self.RESTART_DELAY_SECONDS) + + +# --------------------------------------------------------------------------- +# Module-level helper functions +# --------------------------------------------------------------------------- + +def get_alert_stream_classes() -> list[type[AlertStream]]: + """Return the imported class for each configured alert stream. + + Imports each active stream class from settings.ALERT_STREAMS by its dotted + NAME path. Does NOT instantiate the classes — this is the lightweight + alternative to get_alert_streams() for when you only need access to class + attributes (e.g. STREAM_NAME). + + Streams that are inactive (ACTIVE=False) are skipped. Streams that fail to + import are logged and skipped so a single misconfigured entry doesn't break + the caller. + """ + classes: list[type[AlertStream]] = [] + for stream_config in getattr(settings, 'ALERT_STREAMS', []): + if not stream_config.get('ACTIVE', True): + continue + try: + classes.append(import_string(stream_config['NAME'])) + except (ImportError, AttributeError, KeyError) as exc: + logger.warning( + 'get_alert_stream_classes: could not import %s: %s', + stream_config.get('NAME'), exc, + ) + return classes + + +def get_default_alert_streams() -> list[AlertStream]: + """Return the AlertStream instances configured in settings.ALERT_STREAMS. + + `get_alert_streams()` is the general function. Here, we call that function + and pass in the configuration dictionary from settings.ALERT_STREAMS. + + Raises: + ImproperlyConfigured: if ALERT_STREAMS is not defined in settings, or if + any stream's configuration is invalid. + """ + try: + return get_alert_streams(settings.ALERT_STREAMS) + except AttributeError as err: + raise ImproperlyConfigured( + f'ALERT_STREAMS is not configured in settings.py: {err}' + ) + + +def get_alert_streams(alert_stream_configs: list) -> list[AlertStream]: + """Instantiate and return AlertStream objects from a list of config dicts. + + Use this function if your alert streams are configured somewhere other + than settings.ALERT_STREAMS. + + Each config dict must have: + NAME (str): dotted-path to an AlertStream subclass + OPTIONS (dict): keyword arguments passed to the subclass constructor + ACTIVE (bool, optional): if False, skip this stream (defaults to True) + + Args: + alert_stream_configs: List of configuration dictionaries from ALERT_STREAMS. + + Returns: + A list of instantiated AlertStream subclass objects for active streams. + + Raises: + ImproperlyConfigured: if a NAME cannot be imported, or if Pydantic + validation of OPTIONS fails for any stream. + """ + alert_streams = [] + for alert_stream_config in alert_stream_configs: + if not alert_stream_config.get('ACTIVE', True): + logger.debug( + f'get_alert_streams: skipping inactive stream: {alert_stream_config["NAME"]}' + ) + continue + + # Dynamically import the AlertStream subclass by dotted-path name. + try: + klass = import_string(alert_stream_config['NAME']) + except ImportError as err: + raise ImproperlyConfigured( + f'Could not import AlertStream class "{alert_stream_config["NAME"]}". ' + f'Check the NAME key in your ALERT_STREAMS setting. Error: {err}' + ) + + # Pydantic validates required fields in OPTIONS and raises ValidationError + # with field-level detail if anything is missing or mistyped. + try: + alert_stream: AlertStream = klass(**alert_stream_config.get('OPTIONS', {})) + except ValidationError as err: + raise ImproperlyConfigured( + f'Configuration for {alert_stream_config["NAME"]} is invalid:\n{err}' + ) + + alert_streams.append(alert_stream) + + return alert_streams + + +# --------------------------------------------------------------------------- +# Here's a handler that puts an alert (after normalization) into the the Alerts table +# --------------------------------------------------------------------------- + +def save_alert_to_database(raw_alert: Any, alert_stream: AlertStream, **kwargs: Any) -> Alert | None: + """Persist a raw alert to the database after normalization. + + This could be used as example code for an alert handler that might + handle alerts from multiple streams with different formats. The normalization + step (and injected alert stream) make this possible. + + `alert_stream` is injected by AlertStream.listen() so that this one generic + handler can serve all streams — each stream's normalize_alert() provides + stream-specific field extraction without this function needing to know + which stream it is working with (dependency injection via keyword argument). + + Args: + raw_alert: The stream-specific alert object received from listen(). + alert_stream: The AlertStream instance; provides the normalization function. + **kwargs: Stream-specific extras (e.g., topic, metadata from Hopskotch); + the 'topic' kwarg, if present, is forwarded to normalize_alert() so the + NormalizedAlert.topic field is populated correctly. + + Returns: + The created Alert instance, or None if normalization or save fails. + """ + # get the AlertStream subclass-specific normaization function + normalize = alert_stream.get_normalization_function() + topic: str = kwargs.get('topic', '') + try: + normalized_alert: NormalizedAlert = normalize(raw_alert, topic=topic) + # NormalizedAlert is a Pydantic BaseModel subclass with a model_dump method + alert = Alert.objects.create(**normalized_alert.model_dump()) + return alert + except Exception as ex: + logger.error(f'save_alert_to_database: failed to save alert: {ex}' + f'raw_alert: {raw_alert}') + return None diff --git a/tom_alertstreams/alertstreams/ampel.py b/tom_alertstreams/alertstreams/ampel.py new file mode 100644 index 0000000..9fcdee2 --- /dev/null +++ b/tom_alertstreams/alertstreams/ampel.py @@ -0,0 +1,379 @@ +from __future__ import annotations + +import logging +import random +import time +from datetime import datetime, timezone +from typing import Any, ClassVar, Literal + +# AMPEL imports are commented out because ampel-ztf / ampel-lsst pull in pymongo, +# whose bundled bson package shadows the standalone python-bson required by +# antares_client. Until that conflict is resolved, the real AMPEL classes cannot +# coexist with ANTARES in the same environment. +# +# from ampel.alert.AmpelAlert import AmpelAlert +# from ampel.base.AuxUnitRegister import AuxUnitRegister +# from ampel.lsst.alert.load.KafkaAlertLoader import KafkaAlertLoader +# from ampel.lsst.alert.LSSTAlertSupplier import LSSTAlertSupplier +# from ampel.ztf.alert.ZiAlertSupplier import ZiAlertSupplier +# from ampel.ztf.t0.load.UWAlertLoader import UWAlertLoader + +from tom_alertstreams.alertstreams.alertstream import ( + AlertStream, AlertStreamConfig, NormalizedAlert, _jd_to_datetime, _mjd_to_datetime, +) + +logger = logging.getLogger(__name__) + +# Mock data constants — chosen to be unmistakably non-astronomical: +# (0, 0) is not a real survey pointing; 99.0 is the astronomical sentinel for "no data". +_MOCK_RA = 0.0 +_MOCK_DEC = 0.0 +_MOCK_MAGNITUDE = 99.0 +INTER_ALERT_SLEEP_MIN = 3600 # sixty minutes +INTER_ALERT_SLEEP_MAX = 3600 + + +# --------------------------------------------------------------------------- +# Mock AMPEL stream (for demo use while real AMPEL imports are unavailable) +# --------------------------------------------------------------------------- + +class AmpelMockConfig(AlertStreamConfig): + """Pydantic configuration model for AmpelMockAlertStream. + + Inherits TOPIC_HANDLERS from AlertStreamConfig (a Pydantic BaseModel). + No additional fields needed for mock data generation. + """ + pass + + +class AmpelMockAlertStream(AlertStream): + """Mock AMPEL AlertStream that generates obviously-fake alerts. + + Stands in for the real AMPEL streams while the ampel-ztf / ampel-lsst + packages cannot be installed alongside antares_client (pymongo bson + conflict). Uses the same mock pattern as the other stub streams. + """ + configuration_class = AmpelMockConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'ampel' + IS_MOCK: ClassVar[bool] = True + + def normalize_alert(self, raw_alert: dict, topic: str = '') -> NormalizedAlert: + """Map a mock AMPEL alert dict to a NormalizedAlert. + + Args: + raw_alert: Dict produced by listen(); contains mock field values. + topic: Kafka topic the alert was consumed from. + + Returns: + NormalizedAlert populated from the mock dict fields. + """ + return NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic or raw_alert.get('topic', ''), + observation_time=None, # a mock alert has no real observation + published_time=datetime.fromisoformat(raw_alert['timestamp']), + alert_id=raw_alert['alert_id'], + object_id=raw_alert.get('object_id'), + ra=raw_alert.get('ra'), + dec=raw_alert.get('dec'), + magnitude=raw_alert.get('magnitude'), + raw_payload=raw_alert, + ) + + def listen(self) -> None: + """Generate mock AMPEL alerts and dispatch to configured topic handlers. + + Loops indefinitely, emitting one mock alert per iteration with a random + 6–7 minute delay. Topics are round-robined if multiple are configured. + """ + counter = 0 + topics = list(self.config.TOPIC_HANDLERS.keys()) + + while True: + counter += 1 + topic = topics[counter % len(topics)] + timestamp = datetime.now(timezone.utc) + object_id = f'MOCK-{self.STREAM_NAME.upper()}-{counter:04d}' + alert_id = f'MOCK-{timestamp.strftime("%Y%m%d%H%M%S")}' + mock_alert: dict[str, Any] = { + 'alert_id': alert_id, + 'object_id': object_id, + 'topic': topic, + 'timestamp': timestamp.isoformat(), + 'ra': _MOCK_RA, + 'dec': _MOCK_DEC, + 'magnitude': _MOCK_MAGNITUDE, + 'mock': True, + } + logger.debug(f'AmpelMockAlertStream: mock alert {object_id}') + self.alert_handler[topic](mock_alert, alert_stream=self, topic=topic) + time.sleep(random.uniform(INTER_ALERT_SLEEP_MIN, INTER_ALERT_SLEEP_MAX)) + + +# --------------------------------------------------------------------------- +# ZTF via AMPEL (requires: pip install ampel-ztf) +# --------------------------------------------------------------------------- + +class AmpelZtfConfig(AlertStreamConfig): + """Pydantic configuration for AmpelZtfAlertStream. + + Inherits TOPIC_HANDLERS from AlertStreamConfig (a Pydantic BaseModel). + Fields map to UWAlertLoader parameters. + + Config values: + BOOTSTRAP: Kafka broker address. The UW public broker is the default. + STREAM: 'ztf_uw_public' (all programid=1 alerts) or 'ztf_uw_private' + (adds programid=2). Must match one of the UWAlertLoader literals. + GROUP_NAME: Kafka consumer group prefix. UWAlertLoader appends '-{stream}'. + TIMEOUT: Seconds to wait for messages before the supplier stops iterating. + """ + BOOTSTRAP: str = 'partnership.alerts.ztf.uw.edu:9092' + STREAM: Literal['ztf_uw_private', 'ztf_uw_public'] = 'ztf_uw_public' + GROUP_NAME: str = 'tom-alertstreams' + TIMEOUT: int = 3600 + + +class AmpelZtfAlertStream(AlertStream): + """Consume ZTF alerts from the UW Kafka broker via AMPEL's ZiAlertSupplier. + + Uses ampel-ztf's supplier/loader stack to connect to the University of + Washington's ZTF alert archive. The supplier handles Avro deserialization + and shapes each alert into an AmpelAlert with structured datapoints. + + Alerts are dispatched to the handler registered for the single configured + topic in TOPIC_HANDLERS. ZTF topic metadata is not propagated through the + AMPEL iteration chain, so all alerts use the first (and expected only) + TOPIC_HANDLERS key as the topic. + + Requires: pip install tom-alertstreams[ampel-ztf] + """ + configuration_class = AmpelZtfConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'ampel-ztf' + + def normalize_alert(self, alert: Any, topic: str = '') -> NormalizedAlert: + """Extract common fields from an AMPEL-shaped ZTF alert. + + The AmpelAlert.datapoints[0] is the current candidate (a ReadOnlyDict) + containing jd, ra, dec, magpsf, etc. The original ZTF objectId string + is preserved in AmpelAlert.extra['name'] (alert.stock is an encoded + AMPEL-internal integer, not the human-readable ZTF name). + + Args: + alert: AmpelAlert yielded by ZiAlertSupplier. + topic: Kafka topic name (passed from listen()). + + Returns: + NormalizedAlert with magnitude populated; flux is None. + """ + candidate = alert.datapoints[0] + return NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic, + observation_time=_jd_to_datetime(candidate['jd']), + published_time=None, + alert_id=str(alert.id), + object_id=alert.extra.get('name') if alert.extra else None, + ra=candidate.get('ra'), + dec=candidate.get('dec'), + magnitude=candidate.get('magpsf'), + flux=None, + raw_payload=alert.dict(), + ) + + def listen(self) -> None: + """Connect to the UW Kafka broker and iterate ZTF alerts indefinitely. + + Registers UWAlertLoader with AMPEL's AuxUnitRegister (required for the + supplier to resolve the loader by name), then creates a ZiAlertSupplier + that wraps the loader. UWAlertLoader subscribes to topics in __init__ + (via AllConsumingConsumer), so no context manager is needed. + + Each alert is dispatched to the handler for the single configured topic. + """ + # Deferred import: ampel-ztf's bson (via pymongo) conflicts with + # antares_client's standalone bson. Only import when actually used. + from ampel.alert.AmpelAlert import AmpelAlert # noqa: F811 + from ampel.base.AuxUnitRegister import AuxUnitRegister + from ampel.ztf.alert.ZiAlertSupplier import ZiAlertSupplier + from ampel.ztf.t0.load.UWAlertLoader import UWAlertLoader + + # Register the loader class so AuxUnitRegister can resolve it by name + # when the supplier's UnitModel reference is instantiated. + AuxUnitRegister._dyn['UWAlertLoader'] = UWAlertLoader + + supplier = ZiAlertSupplier( + deserialize='avro', + loader={ + 'unit': 'UWAlertLoader', + 'config': { + 'bootstrap': self.config.BOOTSTRAP, + 'stream': self.config.STREAM, + 'group_name': self.config.GROUP_NAME, + 'timeout': self.config.TIMEOUT, + }, + }, + ) + + # ZTF topics are not propagated through the supplier iteration chain + # (UWAlertLoader.alerts() reads message.topic() for stats but yields + # only the raw bytes). Use the single configured topic key for dispatch. + topic = next(iter(self.config.TOPIC_HANDLERS)) + logger.info( + '%s: listening on %s (stream=%s, group=%s)', + self.STREAM_NAME, self.config.BOOTSTRAP, self.config.STREAM, self.config.GROUP_NAME, + ) + + for alert in supplier: + self.alert_handler[topic](alert, alert_stream=self, topic=topic) + + +# --------------------------------------------------------------------------- +# LSST via AMPEL (requires: pip install ampel-lsst) +# --------------------------------------------------------------------------- + +class AmpelLsstConfig(AlertStreamConfig): + """Pydantic configuration for AmpelLsstAlertStream. + + Inherits TOPIC_HANDLERS from AlertStreamConfig (a Pydantic BaseModel). + Fields map to KafkaAlertLoader / KafkaConsumerBase parameters. + + Config values: + BOOTSTRAP: Kafka broker address (e.g. 'alert-stream-int.lsst.cloud:9094'). + TOPICS: Explicit list of Kafka topics to subscribe to. + GROUP_NAME: Kafka consumer group. LSST brokers require this to be prefixed + with the assigned username (e.g. 'ampel-idfint-tom-alertstreams'). + TIMEOUT: Seconds to wait for messages before the supplier stops iterating. + AVRO_SCHEMA: Schema registry URL for Avro deserialization + (e.g. 'https://alert-schemas-int.lsst.cloud'). None to skip. + SASL_USERNAME / SASL_PASSWORD: SASL/SCRAM credentials assigned by the + Rubin alert stream operators. See DMTN-210 for details. + SASL_MECHANISM: SCRAM variant. Rubin uses SCRAM-SHA-512. + SECURITY_PROTOCOL: Kafka security protocol. Rubin uses SASL_SSL. + KAFKA_CONSUMER_PROPERTIES: Extra confluent_kafka consumer config passed + directly to the underlying DeserializingConsumer. + """ + BOOTSTRAP: str + TOPICS: list[str] + GROUP_NAME: str = 'tom-alertstreams' + TIMEOUT: int = 300 + AVRO_SCHEMA: str | None = None + SASL_USERNAME: str | None = None + SASL_PASSWORD: str | None = None + SASL_MECHANISM: str = 'SCRAM-SHA-512' + SECURITY_PROTOCOL: str = 'SASL_SSL' + KAFKA_CONSUMER_PROPERTIES: dict[str, Any] = {} + + +class AmpelLsstAlertStream(AlertStream): + """Consume LSST alerts from a Kafka broker via AMPEL's LSSTAlertSupplier. + + Uses ampel-lsst's supplier/loader stack to connect to a Rubin-compatible + Kafka broker. The loader handles Avro deserialization (optionally via a + schema registry) and attaches Kafka metadata (__kafka dict) to each alert. + The supplier shapes alerts into AmpelAlert objects with field-name upgrades + (e.g. psFlux → psfFlux, midPointTai → midpointMjdTai, decl → dec). + + Authentication uses SASL_SSL + SCRAM-SHA-512 by default, matching the + Rubin alert distribution system (DMTN-210). + + Requires: pip install tom-alertstreams[ampel-lsst] + """ + configuration_class = AmpelLsstConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'ampel-lsst' + + def normalize_alert(self, alert: Any, topic: str = '') -> NormalizedAlert: + """Extract common fields from an AMPEL-shaped LSST alert. + + AmpelAlert.datapoints[0] is the triggering diaSource (a ReadOnlyDict) + with field-upgraded names: midpointMjdTai, psfFlux, ra, dec, band, etc. + AmpelAlert.stock is the diaObjectId. Kafka topic metadata is available + in alert.extra['kafka']['topic'] (preserved by KafkaAlertLoader). + + Args: + alert: AmpelAlert yielded by LSSTAlertSupplier. + topic: Kafka topic name (passed from listen()). + + Returns: + NormalizedAlert with flux populated (nanojansky); magnitude is None. + """ + dia_source = alert.datapoints[0] + # Topic from Kafka metadata, falling back to the caller-provided value. + kafka_topic = '' + if alert.extra and 'kafka' in alert.extra: + kafka_topic = alert.extra['kafka'].get('topic', '') + return NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic or kafka_topic, + observation_time=_mjd_to_datetime(dia_source['midpointMjdTai']), + published_time=None, + alert_id=str(alert.id), + object_id=str(alert.stock), + ra=dia_source.get('ra'), + dec=dia_source.get('dec'), + magnitude=None, + flux=dia_source.get('psfFlux'), + raw_payload=alert.dict(), + ) + + def listen(self) -> None: + """Connect to a Kafka broker and iterate LSST alerts indefinitely. + + Registers KafkaAlertLoader with AMPEL's AuxUnitRegister, then creates + an LSSTAlertSupplier. Unlike ZTF, the LSST loader subscribes to topics + in __enter__() (a context manager), so we wrap the iteration in a `with` + block on the loader. + + SASL credentials are passed to the loader via kafka_consumer_properties + rather than AMPEL's NamedSecret-based SASLAuthentication, which requires + AMPEL's secret store infrastructure. Passing them as plain consumer + properties is simpler and sufficient for our use case. + """ + # Deferred import: ampel-lsst's bson (via pymongo) conflicts with + # antares_client's standalone bson. Only import when actually used. + from ampel.base.AuxUnitRegister import AuxUnitRegister + from ampel.lsst.alert.load.KafkaAlertLoader import KafkaAlertLoader + from ampel.lsst.alert.LSSTAlertSupplier import LSSTAlertSupplier + + AuxUnitRegister._dyn['KafkaAlertLoader'] = KafkaAlertLoader + + # Build loader config, merging SASL auth into kafka_consumer_properties. + kafka_props = dict(self.config.KAFKA_CONSUMER_PROPERTIES) + if self.config.SASL_USERNAME: + kafka_props.update({ + 'security.protocol': self.config.SECURITY_PROTOCOL, + 'sasl.mechanism': self.config.SASL_MECHANISM, + 'sasl.username': self.config.SASL_USERNAME, + 'sasl.password': self.config.SASL_PASSWORD, + }) + + loader_config: dict[str, Any] = { + 'bootstrap': self.config.BOOTSTRAP, + 'topics': self.config.TOPICS, + 'group_name': self.config.GROUP_NAME, + 'timeout': self.config.TIMEOUT, + 'kafka_consumer_properties': kafka_props, + } + if self.config.AVRO_SCHEMA: + loader_config['avro_schema'] = self.config.AVRO_SCHEMA + + supplier = LSSTAlertSupplier( + deserialize=None, # KafkaAlertLoader handles deserialization + loader={'unit': 'KafkaAlertLoader', 'config': loader_config}, + ) + + logger.info( + '%s: listening on %s (topics=%s, group=%s)', + self.STREAM_NAME, self.config.BOOTSTRAP, self.config.TOPICS, self.config.GROUP_NAME, + ) + + # KafkaAlertLoader subscribes in __enter__() — must use context manager. + with supplier.alert_loader: + fallback_topic = next(iter(self.config.TOPIC_HANDLERS)) + for alert in supplier: + # Extract real topic from Kafka metadata preserved by KafkaAlertLoader. + topic = '' + if alert.extra and 'kafka' in alert.extra: + topic = alert.extra['kafka'].get('topic', '') + if not topic: + topic = fallback_topic + self.alert_handler[topic](alert, alert_stream=self, topic=topic) diff --git a/tom_alertstreams/alertstreams/antares.py b/tom_alertstreams/alertstreams/antares.py index bd17fa4..6c07331 100644 --- a/tom_alertstreams/alertstreams/antares.py +++ b/tom_alertstreams/alertstreams/antares.py @@ -1,30 +1,283 @@ +from __future__ import annotations + import logging -from .alertstream import AlertStream +import random +import time +from datetime import datetime, timezone +from typing import Any, ClassVar + from antares_client.stream import StreamingClient +from tom_alertstreams.alertstreams.alertstream import ( + AlertStream, AlertStreamConfig, NormalizedAlert, _mjd_to_datetime, +) + logger = logging.getLogger(__name__) +# Mock data constants — chosen to be unmistakably non-astronomical: +# (0, 0) is not a real survey pointing; 99.0 is the astronomical sentinel for "no data". +_MOCK_RA = 0.0 +_MOCK_DEC = 0.0 +_MOCK_MAGNITUDE = 99.0 +INTER_ALERT_SLEEP_MIN = 360 # six minutes +INTER_ALERT_SLEEP_MAX = 420 # seven minutes -class AntaresAlertStream(AlertStream): + +# --------------------------------------------------------------------------- +# Pydantic configuration models +# --------------------------------------------------------------------------- + +class AntaresConfig(AlertStreamConfig): + """Pydantic configuration model for AntaresAlertStream. + + Inherits from AlertStreamConfig (a Pydantic BaseModel), so Pydantic validates + that API_KEY and API_SECRET are present and raises descriptive errors if not. + + Fields: + API_KEY: ANTARES API key (required). Obtain at https://antares.noirlab.edu. + API_SECRET: ANTARES API secret (required). Obtain at https://antares.noirlab.edu. + TOPIC_HANDLERS: Inherited from AlertStreamConfig. Maps ANTARES filter topic + names to handler dotted-paths. ANTARES topics are filter outputs, not + survey-specific — a single topic can carry both ZTF and LSST loci. + Known topics: 'extragalactic_staging', 'nuclear_transient_staging', + 'in_shadow_virgo'. See the ANTARES tags page for the full list. + GROUP: Kafka consumer group ID. Distinct group IDs let multiple TOM instances + consume the same stream independently. + SSL_CA_LOCATION: Path to a TLS Certificate Authority (CA) certificate file. + The antares_client package bundles a default CA cert that works for the + current broker. Set this only if a different CA cert is needed. + When None (the default), the bundled cert is used. + ENABLE_AUTO_COMMIT: Whether Kafka should auto-commit offsets. Set to False + for at-least-once processing with manual offset management. """ - Wrapper for the ANTARES broker streaming client. See https://nsf-noirlab.gitlab.io/csdc/antares/client/. + API_KEY: str + API_SECRET: str + GROUP: str = 'tom-alertstreams' + SSL_CA_LOCATION: str | None = None + ENABLE_AUTO_COMMIT: bool = True + + +class AntaresMockConfig(AlertStreamConfig): + """Pydantic configuration model for AntaresMockAlertStream. + + Inherits TOPIC_HANDLERS from AlertStreamConfig (a Pydantic BaseModel). + No additional fields needed for mock data generation. """ - required_keys = ['API_KEY', 'API_SECRET', 'TOPIC_HANDLERS'] - allowed_keys = ['API_KEY', 'API_SECRET', 'TOPIC_HANDLERS', 'GROUP', 'SSL_CA_LOCATION', 'ENABLE_AUTO_COMMIT'] - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - logger.debug(f'AntaresAlertStream.__init__() kwargs: {kwargs}') - optional_keys = set(self.allowed_keys) - set(self.required_keys) - options = {key.lower(): kwargs[key] for key in optional_keys if key in kwargs} - self.stream = StreamingClient( - topics=self.topic_handlers.keys(), - api_key=self.api_key, - api_secret=self.api_secret, - **options + pass + + +# --------------------------------------------------------------------------- +# Mock ANTARES stream (for demo use without API credentials) +# --------------------------------------------------------------------------- + +class AntaresMockAlertStream(AlertStream): + """Mock ANTARES AlertStream that generates obviously-fake alerts. + + Provides a demo fallback when ANTARES API credentials are not available. + Uses the same mock pattern as the other stub streams (sentinel coordinates, + sentinel magnitude, 6–7 minute sleep between alerts). + + The mock stream uses STREAM_NAME='antares' so it occupies the same dashboard + slot as the real ANTARES streams would if they were the only ones configured. + """ + configuration_class = AntaresMockConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'antares' + IS_MOCK: ClassVar[bool] = True + + def normalize_alert(self, raw_alert: dict, topic: str = '') -> NormalizedAlert: + """Map a mock ANTARES alert dict to a NormalizedAlert. + + Args: + raw_alert: Dict produced by listen(); contains mock field values. + topic: Topic the alert was generated for. + + Returns: + NormalizedAlert populated from the mock dict fields. + """ + return NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic or raw_alert.get('topic', ''), + observation_time=None, # a mock alert has no real observation + published_time=datetime.fromisoformat(raw_alert['timestamp']), + alert_id=raw_alert['alert_id'], + object_id=raw_alert.get('object_id'), + ra=raw_alert.get('ra'), + dec=raw_alert.get('dec'), + magnitude=raw_alert.get('magnitude'), + raw_payload=raw_alert, ) - def listen(self): - for topic, locus in self.stream.iter(): - base_topic = topic.removeprefix(self.stream._TOPIC_PREFIX) - logger.info(f"received {locus.locus_id} on {base_topic}") - self.alert_handler[base_topic](locus) + def listen(self) -> None: + """Generate mock ANTARES alerts and dispatch to configured topic handlers. + + Loops indefinitely, emitting one mock alert per iteration with a random + 6–7 minute delay. Topics are round-robined if multiple are configured. + """ + counter = 0 + topics = list(self.config.TOPIC_HANDLERS.keys()) + + while True: + counter += 1 + topic = topics[counter % len(topics)] + timestamp = datetime.now(timezone.utc) + object_id = f'MOCK-{self.STREAM_NAME.upper()}-{counter:04d}' + alert_id = f'MOCK-{timestamp.strftime("%Y%m%d%H%M%S")}' + mock_alert: dict[str, Any] = { + 'alert_id': alert_id, + 'object_id': object_id, + 'topic': topic, + 'timestamp': timestamp.isoformat(), + 'ra': _MOCK_RA, + 'dec': _MOCK_DEC, + 'magnitude': _MOCK_MAGNITUDE, + 'mock': True, + } + logger.debug(f'AntaresMockAlertStream: mock alert {object_id}') + self.alert_handler[topic](mock_alert, alert_stream=self, topic=topic) + time.sleep(random.uniform(INTER_ALERT_SLEEP_MIN, INTER_ALERT_SLEEP_MAX)) + + +# --------------------------------------------------------------------------- +# Real ANTARES stream — unified for all surveys (ZTF, LSST, etc.) +# --------------------------------------------------------------------------- + +class AntaresAlertStream(AlertStream): + """ANTARES alert stream for ZTF, LSST, and future survey alerts. + + ANTARES Kafka topics are filter outputs, not survey-specific — a single topic + like 'extragalactic_staging' can carry loci with ZTF data, LSST data, or both. + This class handles all surveys through a single normalize_alert() that extracts + whatever survey data is present on each locus. + + The StreamingClient is created inside listen() using a context manager (not in + __init__) so the Kafka consumer is properly closed on errors or shutdown. + + Known topics (from the ANTARES tags page at https://antares.noirlab.edu): + ZTF: 'extragalactic_staging', 'nuclear_transient_staging' + Mixed/LSST: 'in_shadow_virgo' + Topic names generally follow the pattern '{tag_name}_staging', though some + (like 'in_shadow_virgo') omit the suffix. + + Configuration example (settings.py ALERT_STREAMS entry):: + + { + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.antares.AntaresAlertStream', + 'OPTIONS': { + 'API_KEY': os.environ.get('ANTARES_API_KEY', ''), + 'API_SECRET': os.environ.get('ANTARES_API_SECRET', ''), + 'TOPIC_HANDLERS': { + # ZTF topics + 'extragalactic_staging': 'tom_alertstreams.alertstreams.alertstream.save_alert_to_database', + 'nuclear_transient_staging': 'tom_alertstreams.alertstreams.alertstream.save_alert_to_database', + # LSST topics + 'in_shadow_virgo': 'tom_alertstreams.alertstreams.alertstream.save_alert_to_database', + }, + }, + } + """ + configuration_class = AntaresConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'antares' + + def listen(self) -> None: + """Consume ANTARES loci and dispatch to configured topic handlers. + + Creates a StreamingClient as a context manager so the underlying Kafka + consumer is properly closed on errors, KeyboardInterrupt, or normal exit. + + ANTARES delivers (topic, locus) pairs via StreamingClient.iter(). The topic + includes the stream prefix (e.g. 'client.extragalactic_staging'); we strip + the prefix to match the base topic names used in TOPIC_HANDLERS. + """ + # Build optional kwargs from config — only pass non-default values so the + # StreamingClient uses its own defaults for unset fields. + optional_kwargs: dict[str, Any] = {} + if self.config.SSL_CA_LOCATION is not None: + optional_kwargs['ssl_ca_location'] = self.config.SSL_CA_LOCATION + if not self.config.ENABLE_AUTO_COMMIT: + optional_kwargs['enable_auto_commit'] = self.config.ENABLE_AUTO_COMMIT + + with StreamingClient( + topics=list(self.config.TOPIC_HANDLERS.keys()), + api_key=self.config.API_KEY, + api_secret=self.config.API_SECRET, + group=self.config.GROUP, + **optional_kwargs, + ) as antares_streaming_client: + for topic, locus in antares_streaming_client.iter(): + base_topic = topic.removeprefix(antares_streaming_client._TOPIC_PREFIX) + logger.info(f'{self.STREAM_NAME} received {locus.locus_id} on {base_topic}') + self.alert_handler[base_topic](locus, alert_stream=self, topic=base_topic) + + def normalize_alert(self, raw_alert: Any, topic: str = '') -> NormalizedAlert: + """Extract common fields from an ANTARES Locus object. + + Handles both ZTF and LSST data by extracting whatever survey-specific + properties are present on the locus. Field mappings discovered via + ex_antares_ztf.py and ex_antares_lsst.py introspection: + + - timestamp: properties['newest_alert_observation_time'] (MJD float) + - magnitude: properties['newest_alert_magnitude'] (ANTARES-enriched) + - object_id: LSST dia_object_id (nested) → ZTF ztf_object_id (flat) → locus_id + - alert_id: locus.locus_id (used by AntaresPresenter for locus page URL) + + The full Locus object is not JSON-serializable, so raw_payload is left empty. + Handlers that need Locus-specific data should access raw_alert directly. + + Args: + raw_alert: An antares_client.models.Locus object. + topic: The ANTARES topic (e.g. 'extragalactic_staging'). + + Returns: + NormalizedAlert with survey-appropriate fields populated. + """ + alert_properties = raw_alert.properties or {} + + # Extract properties from the raw_alert for transfer to NormalizedAlert. + + # Timestamp from ANTARES-enriched locus property (MJD float). + # This avoids lazy-loading locus.alerts, which triggers an HTTP API call. + mjd = alert_properties.get('newest_alert_observation_time') + timestamp = _mjd_to_datetime(mjd) if mjd is not None else None + + # Magnitude — ANTARES-enriched, populated for ZTF loci. + magnitude = alert_properties.get('newest_alert_magnitude') + + # Object ID — try LSST nested structure first, then ZTF flat property, then locus_id. + # ANTARES stores LSST IDs in a nested dict: properties['survey']['lsst']['dia_object_id'] → list + # ZTF IDs are a flat property: properties['ztf_object_id'] → str + object_id = raw_alert.locus_id # fallback + survey = alert_properties.get('survey', {}) + lsst_survey = survey.get('lsst', {}) if isinstance(survey, dict) else {} + dia_object_ids = lsst_survey.get('dia_object_id', []) + if dia_object_ids: + object_id = dia_object_ids[0] + elif alert_properties.get('ztf_object_id'): + object_id = alert_properties['ztf_object_id'] + + # Flux — LSST uses flux (nanojansky) instead of magnitude. + # Property key TBD: no LSST-only locus observed from the stream yet. + # TODO: once the flux property key is known, extract it here. + flux = None + + # Log LSST locus properties so we can discover the flux key from production data. + # Remove this logging once we've confirmed the flux property key and updated + # the extraction above. + if dia_object_ids: + logger.info(f'LSST locus detected: {raw_alert.locus_id} — ' + f'full properties: {alert_properties}') + + normalized_alert = NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic, + observation_time=timestamp, + published_time=None, # ANTARES locus exposes no broker-publish time we extract + alert_id=str(raw_alert.locus_id), + object_id=str(object_id), + ra=float(raw_alert.ra) if raw_alert.ra is not None else None, + dec=float(raw_alert.dec) if raw_alert.dec is not None else None, + magnitude=float(magnitude) if magnitude is not None else None, + flux=flux, + raw_payload={}, + ) + return normalized_alert diff --git a/tom_alertstreams/alertstreams/babamul.py b/tom_alertstreams/alertstreams/babamul.py new file mode 100644 index 0000000..14f0dfc --- /dev/null +++ b/tom_alertstreams/alertstreams/babamul.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +import logging +from typing import ClassVar, Literal + +from babamul import AlertConsumer, LsstAlert, ZtfAlert +from pydantic import Field + +from tom_alertstreams.alertstreams.alertstream import AlertStream, AlertStreamConfig, NormalizedAlert + +logger = logging.getLogger(__name__) + + +class BabamulConfig(AlertStreamConfig): + """Pydantic configuration model for BabamulAlertStream. + + Inherits TOPIC_HANDLERS from AlertStreamConfig (a Pydantic BaseModel). + + Fields: + BABAMUL_KAFKA_USERNAME: Kafka username for the Babamul broker (required). + BABAMUL_KAFKA_PASSWORD: Kafka password for the Babamul broker (required). + BABAMUL_GROUP_ID: Kafka consumer group ID. Alerts are partitioned across + consumers sharing the same group_id, so each TOM instance should use + a unique group_id to receive all alerts. + BABAMUL_AUTO_COMMIT: Whether to auto-commit Kafka offsets after consuming. + False (default) means offsets are not committed, so restarting the + consumer replays from the configured BABAMUL_OFFSET position. + BABAMUL_OFFSET: Where to start reading when no committed offset exists. + 'EARLIEST' replays all available alerts; 'LATEST' starts from new ones. + """ + BABAMUL_KAFKA_USERNAME: str = Field(min_length=1) # don't accept an empty string + BABAMUL_KAFKA_PASSWORD: str = Field(min_length=1) + BABAMUL_GROUP_ID: str = Field(min_length=1) + BABAMUL_AUTO_COMMIT: bool = False + BABAMUL_OFFSET: Literal['earliest', 'latest'] = 'latest' + + +class BabamulAlertStream(AlertStream): + """AlertStream implementation for Babamul (https://github.com/boom-astro/babamul). + """ + configuration_class = BabamulConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'babamul' + + def normalize_alert(self, raw_alert: ZtfAlert | LsstAlert, topic: str = '') -> NormalizedAlert: + """Convert a babamul alert to a NormalizedAlert. + + ZtfAlert and LsstAlert are Pydantic BaseModel subclasses provided by the + ``babamul`` package. Because they are Pydantic models we can: + - access fields via typed attributes (raw_alert.objectId, candidate.ra, etc.) + - serialize to JSON-safe dicts with model_dump(mode='json'), which handles + datetimes, enums (e.g. Band), and nested models automatically + - skip defensive getattr() / try-except — field access is guaranteed by the schema + + Args: + raw_alert: A babamul ZtfAlert or LsstAlert from AlertConsumer. + topic: The Kafka topic the alert arrived on. Falls back to + raw_alert.topic if not provided. + + Returns: + NormalizedAlert with fields extracted from the babamul alert's candidate. + """ + candidate = raw_alert.candidate + + return NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic or raw_alert.topic or '', + observation_time=candidate.datetime, + published_time=None, + alert_id=str(raw_alert.candid), + object_id=raw_alert.objectId, + ra=candidate.ra, + dec=candidate.dec, + magnitude=candidate.magpsf, + raw_payload=raw_alert.model_dump(mode='json'), + ) + + def listen(self) -> None: + """Consume Babamul alerts and dispatch to configured topic handlers. + + Opens a babamul AlertConsumer as a context manager and iterates over + incoming alerts. Each alert is dispatched to the handler for its topic. + + babamul raises on a dropped/unreachable broker (BabamulConnectionError), + a SASL failure (AuthenticationError), or a bad message (DeserializationError) + — all from inside the `for alert in consumer:` iteration. We let those + propagate: AlertStream.run() supervises listen() and reconnects after a + backoff, so connection resilience lives once in the base class rather than + being re-implemented per broker. + """ + topics = list(self.config.TOPIC_HANDLERS.keys()) + + with AlertConsumer( + topics=topics, + username=self.config.BABAMUL_KAFKA_USERNAME, + password=self.config.BABAMUL_KAFKA_PASSWORD, + group_id=self.config.BABAMUL_GROUP_ID, + offset=self.config.BABAMUL_OFFSET, + auto_commit=self.config.BABAMUL_AUTO_COMMIT, + ) as consumer: + alert: ZtfAlert | LsstAlert + for alert in consumer: # yields ZtfAlert | LsstAlert (Pydantic models) + topic = alert.topic or '' + if topic not in self.alert_handler: + logger.warning( + f'BabamulAlertStream: alert from topic "{topic}" has no handler. ' + f'Configured topics: {list(self.alert_handler.keys())}' + ) + continue + + logger.debug(f'BabamulAlertStream: alert {alert.objectId} (candid={alert.candid}) on {topic}') + self.alert_handler[topic](alert, alert_stream=self, topic=topic) diff --git a/tom_alertstreams/alertstreams/fink.py b/tom_alertstreams/alertstreams/fink.py new file mode 100644 index 0000000..6e92093 --- /dev/null +++ b/tom_alertstreams/alertstreams/fink.py @@ -0,0 +1,326 @@ +from __future__ import annotations + +import logging +import random +import time +import warnings +from datetime import datetime, timezone +from typing import Any, ClassVar, Literal + +from fink_client.consumer import AlertConsumer, extract_id_from_lsst +from pydantic import model_validator + +from tom_alertstreams.alertstreams.alertstream import ( + AlertStream, AlertStreamConfig, NormalizedAlert, _jd_to_datetime, _mjd_to_datetime, +) + +logger = logging.getLogger(__name__) + +# Mock data constants — chosen to be unmistakably non-astronomical: +# (0, 0) is not a real survey pointing; 99.0 is the astronomical sentinel for "no data". +_MOCK_RA = 0.0 +_MOCK_DEC = 0.0 +_MOCK_MAGNITUDE = 99.0 +INTER_ALERT_SLEEP_MIN = 360 # six minutes +INTER_ALERT_SLEEP_MAX = 420 # seven minutes + + +# --------------------------------------------------------------------------- +# Pydantic configuration models +# --------------------------------------------------------------------------- + +class FinkMockConfig(AlertStreamConfig): + """Pydantic configuration model for FinkMockAlertStream. + + Inherits TOPIC_HANDLERS from AlertStreamConfig (a Pydantic BaseModel). + No additional fields needed for mock data generation. + """ + pass + + +class FinkConfig(AlertStreamConfig): + """Pydantic configuration for FinkAlertStream. + + Fields: + FINK_USERNAME: Fink username (required). Obtained via Fink registration. + FINK_PASSWORD: Fink password (optional — not all accounts require one). + FINK_GROUP_ID: Kafka consumer group ID (required). + FINK_SERVER: Kafka broker address for the chosen survey (required). + ZTF: 'kafka-ztf.fink-broker.org:24499' + LSST: 'kafka-lsst.fink-broker.org:24499' + FINK_SURVEY: Survey name: 'ztf' or 'lsst'. Default: 'lsst'. + TOPIC_HANDLERS: Inherited from AlertStreamConfig. Maps Fink topic names + to handler dotted-paths. Topics are science-filter outputs specific + to each survey (e.g. 'fink_sn_candidates_ztf' for ZTF). + """ + FINK_USERNAME: str + FINK_PASSWORD: str | None = None + FINK_GROUP_ID: str + FINK_SERVER: str + FINK_SURVEY: Literal['ztf', 'lsst'] = 'lsst' + + @model_validator(mode='after') + def _check_server_matches_survey(self) -> FinkConfig: + """Warn if FINK_SERVER doesn't contain the survey name. + + Fink uses separate Kafka brokers per survey. A mismatch between + FINK_SURVEY and FINK_SERVER likely means the wrong broker address. + """ + if self.FINK_SURVEY not in self.FINK_SERVER: + warnings.warn( + f'FINK_SERVER ({self.FINK_SERVER}) does not contain ' + f'survey name ({self.FINK_SURVEY}). Please verify that the ' + f'broker address matches the chosen survey.', + stacklevel=2, + ) + return self + + @model_validator(mode='after') + def _check_topics_match_survey(self) -> FinkConfig: + """Fail loudly if a topic is tagged for the wrong survey. + + Fink science topics follow the '_' naming convention + (e.g. 'fink_sn_candidates_ztf', 'fink_sn_candidates_lsst'). A topic + ending in the opposite survey's suffix is unambiguously a misconfiguration: + normalize_alert() branches on FINK_SURVEY, so a ZTF topic consumed under + FINK_SURVEY='lsst' would hit the LSST branch, raise KeyError on the missing + 'diaObject' key, get swallowed by save_alert_to_database()'s broad except, + and silently drop every alert. + + Unlike the server check (a substring heuristic that only warns), this is a + hard error: the suffix mismatch is definitive, so we raise at config-load + time. get_alert_streams() converts the resulting ValidationError into an + ImproperlyConfigured, so the misconfiguration surfaces at startup rather + than as silent data loss at runtime. + + Topics with no survey suffix (test topics, the mock's 'fink.test') are + never flagged — only a topic carrying the *opposite* survey's suffix is. + """ + opposite_survey = 'ztf' if self.FINK_SURVEY == 'lsst' else 'lsst' + mismatched = [topic for topic in self.TOPIC_HANDLERS if topic.endswith(f'_{opposite_survey}')] + if mismatched: + raise ValueError( + f"FINK_SURVEY is '{self.FINK_SURVEY}' but these topics are tagged for " + f"'{opposite_survey}': {mismatched}. Fink topics follow the " + f"'_' convention; a survey/topic mismatch routes alerts " + f"to the wrong normalize_alert() branch and silently drops them." + ) + return self + + +# --------------------------------------------------------------------------- +# Mock Fink stream (for demo use without Fink credentials) +# --------------------------------------------------------------------------- + +class FinkMockAlertStream(AlertStream): + """Mock Fink AlertStream that generates obviously-fake alerts. + + Provides a demo fallback when Fink credentials are not available. + Uses the same mock pattern as the other stub streams (sentinel coordinates, + sentinel magnitude, 6–7 minute sleep between alerts). + + The mock stream uses STREAM_NAME='fink' so it occupies the same dashboard + slot as the real Fink stream would if it were the only one configured. + """ + configuration_class = FinkMockConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'fink' + IS_MOCK: ClassVar[bool] = True + + def normalize_alert(self, raw_alert: dict, topic: str = '') -> NormalizedAlert: + """Map a mock Fink alert dict to a NormalizedAlert. + + Args: + raw_alert: Dict produced by listen(); contains mock field values. + topic: Topic the alert was generated for. + + Returns: + NormalizedAlert populated from the mock dict fields. + """ + return NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic or raw_alert.get('topic', ''), + observation_time=None, # a mock alert has no real observation + published_time=datetime.fromisoformat(raw_alert['timestamp']), + alert_id=raw_alert['alert_id'], + object_id=raw_alert.get('object_id'), + ra=raw_alert.get('ra'), + dec=raw_alert.get('dec'), + magnitude=raw_alert.get('magnitude'), + raw_payload=raw_alert, + ) + + def listen(self) -> None: + """Generate mock Fink alerts and dispatch to configured topic handlers. + + Loops indefinitely, emitting one mock alert per iteration with a random + 6–7 minute delay. Topics are round-robined if multiple are configured. + """ + counter = 0 + topics = list(self.config.TOPIC_HANDLERS.keys()) + + while True: + counter += 1 + topic = topics[counter % len(topics)] + timestamp = datetime.now(timezone.utc) + object_id = f'MOCK-{self.STREAM_NAME.upper()}-{counter:04d}' + alert_id = f'MOCK-{timestamp.strftime("%Y%m%d%H%M%S")}' + mock_alert: dict[str, Any] = { + 'alert_id': alert_id, + 'object_id': object_id, + 'topic': topic, + 'timestamp': timestamp.isoformat(), + 'ra': _MOCK_RA, + 'dec': _MOCK_DEC, + 'magnitude': _MOCK_MAGNITUDE, + 'mock': True, + } + logger.debug(f'FinkMockAlertStream: mock alert {object_id}') + self.alert_handler[topic](mock_alert, alert_stream=self, topic=topic) + time.sleep(random.uniform(INTER_ALERT_SLEEP_MIN, INTER_ALERT_SLEEP_MAX)) + + +# --------------------------------------------------------------------------- +# Real Fink stream — supports both ZTF and LSST surveys +# --------------------------------------------------------------------------- + +class FinkAlertStream(AlertStream): + """Fink alert stream for ZTF and LSST alerts. + + Uses fink-client's AlertConsumer to receive Avro-encoded alerts from + the Fink broker. The FINK_SURVEY config field determines which alert + schema to expect (ZTF candidate-based vs LSST diaSource-based). + + Each FinkAlertStream instance reads from a single survey. To read both + ZTF and LSST, configure two entries in ALERT_STREAMS with different + FINK_SURVEY and FINK_SERVER values. + + Configuration example (settings.py ALERT_STREAMS entry):: + + { + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.fink.FinkAlertStream', + 'OPTIONS': { + 'FINK_USERNAME': os.environ.get('FINK_USERNAME', ''), + 'FINK_GROUP_ID': os.environ.get('FINK_GROUPID', ''), + 'FINK_SERVER': os.environ.get('FINK_LSST_SERVER', ''), + 'FINK_SURVEY': 'lsst', + 'TOPIC_HANDLERS': { + 'fink_sn_candidates_lsst': 'tom_alertstreams.alertstreams.alertstream.save_alert_to_database', + }, + }, + } + """ + configuration_class = FinkConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'fink' + + def listen(self) -> None: + """Consume Fink alerts and dispatch to configured topic handlers. + + Creates an AlertConsumer as a context manager so the underlying Kafka + consumer is properly closed on errors, KeyboardInterrupt, or normal exit. + """ + # Build the consumer config dict directly from Django settings. + # fink-client's README describes a CLI workflow (fink_client_register writes + # credentials to ~/.finkclient/{survey}_credentials.yml, then fink_consumer + # reads them via load_credentials()). We bypass that entirely — AlertConsumer + # accepts a plain config dict, so we read from env vars via Django settings + # instead. This avoids needing CLI setup in deployment (Docker, etc.). + # + # 'bootstrap.servers' is Kafka's term for the initial broker address used + # to discover the cluster — we expose it as FINK_SERVER in our config. + consumer_config: dict[str, str] = { + 'username': self.config.FINK_USERNAME, + 'group.id': self.config.FINK_GROUP_ID, + 'bootstrap.servers': self.config.FINK_SERVER, + } + if self.config.FINK_PASSWORD is not None: + consumer_config['password'] = self.config.FINK_PASSWORD + + with AlertConsumer( + topics=list(self.config.TOPIC_HANDLERS.keys()), + config=consumer_config, + survey=self.config.FINK_SURVEY, + ) as fink_consumer: + while True: + topic, alert, key = fink_consumer.poll(timeout=30) + if topic is None: + continue # timeout with no message, retry + logger.info(f'{self.STREAM_NAME} received alert on {topic}') + self.alert_handler[topic](alert, alert_stream=self, topic=topic) + + def normalize_alert(self, raw_alert: dict, topic: str = '') -> NormalizedAlert: + """Extract common fields from a Fink alert dict. + + Handles both ZTF and LSST alert schemas based on the FINK_SURVEY config. + Fink alerts are plain dicts (decoded from Avro), so raw_payload is populated. + + ZTF alert structure: + - objectId, candid at top level + - candidate dict: jd, ra, dec, magpsf, fid, ... + LSST alert structure: + - diaSource dict: diaSourceId, midpointMjdTai, ra, dec, psFlux, ... + - diaObject dict: diaObjectId (or mpc_orbits.designation for SSOs) + + Args: + raw_alert: Alert dict decoded from Fink's Avro stream. + topic: The Fink topic the alert was consumed from. + + Returns: + NormalizedAlert with survey-appropriate fields populated. + """ + # Strip cutout stamp data from raw_payload — these are compressed FITS images + # (cutoutScience, cutoutTemplate, cutoutDifference) that are large, binary, and + # not JSON-serializable. Keep everything else for debugging and data exploration. + raw_payload = {k: v for k, v in raw_alert.items() if not k.startswith('cutout')} + + # Fink stamps its own processing times. brokerEndProcessTimestamp is when Fink + # finished processing and published this alert — a naive-UTC ISO string. The gap + # between it and observation_time is Fink's processing latency. + published_str = raw_alert.get('brokerEndProcessTimestamp') + published_time = ( + datetime.fromisoformat(published_str).replace(tzinfo=timezone.utc) + if published_str else None + ) + + if self.config.FINK_SURVEY == 'ztf': + # ZTF alert structure: top-level objectId/candid, nested candidate dict + candidate = raw_alert.get('candidate', {}) + jd = candidate.get('jd') + timestamp = _jd_to_datetime(jd) if jd is not None else None + normalized_alert = NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic, + observation_time=timestamp, + published_time=published_time, + alert_id=str(raw_alert.get('candid', '')), + object_id=raw_alert.get('objectId'), + ra=candidate.get('ra'), + dec=candidate.get('dec'), + magnitude=candidate.get('magpsf'), + flux=None, + raw_payload=raw_payload, + ) + else: + # LSST alert structure: nested diaSource/diaObject dicts + dia_source = raw_alert.get('diaSource', {}) + mjd_tai = dia_source.get('midpointMjdTai') + timestamp = _mjd_to_datetime(mjd_tai) if mjd_tai is not None else None + + # extract_id_from_lsst (from fink-client) handles static objects + # (diaObject.diaObjectId) vs moving objects (mpc_orbits.designation). + object_id, _ = extract_id_from_lsst(raw_alert) + + normalized_alert = NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic, + observation_time=timestamp, + published_time=published_time, + alert_id=str(dia_source.get('diaSourceId', '')), + object_id=str(object_id), + ra=dia_source.get('ra'), + dec=dia_source.get('dec'), + magnitude=None, + flux=dia_source.get('psFlux'), + raw_payload=raw_payload, + ) + return normalized_alert diff --git a/tom_alertstreams/alertstreams/gcn.py b/tom_alertstreams/alertstreams/gcn.py index b6f32ef..f55562e 100644 --- a/tom_alertstreams/alertstreams/gcn.py +++ b/tom_alertstreams/alertstreams/gcn.py @@ -1,69 +1,211 @@ +from __future__ import annotations + +import json import logging +from datetime import datetime, timezone +from typing import Any, ClassVar from gcn_kafka import Consumer -from tom_alertstreams.alertstreams.alertstream import AlertStream +from pydantic import Field +from tom_alertstreams.alertstreams.alertstream import ( + AlertStream, AlertStreamConfig, NormalizedAlert, is_in_hourly_window, save_alert_to_database, +) logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) -class GCNClassicAlertStream(AlertStream): - """ +class GCNConfig(AlertStreamConfig): + """Pydantic configuration model for GCNClassicAlertStream. - Pre-requisite: visit gcn.nasa.gov and sign-up to get your client_id and - client_secret. + Inherits from AlertStreamConfig (a Pydantic BaseModel), so Pydantic validates + that GCN_KAFKA_CLIENT_ID and GCN_KAFKA_CLIENT_SECRET are present and non-empty. + + Fields: + GCN_KAFKA_CLIENT_ID: GCN client ID (required). Register at https://gcn.nasa.gov/quickstart. + GCN_KAFKA_CLIENT_SECRET: GCN client secret (required). Register at https://gcn.nasa.gov/quickstart. + TOPIC_HANDLERS: Inherited from AlertStreamConfig. Maps GCN topic names to + handler dotted-paths. Example topic: 'gcn.circulars'. + DOMAIN: Kafka broker domain. Defaults to the public GCN broker. + KAFKA_CONFIG: Optional dict passed directly to the underlying Confluent Kafka + Consumer for advanced configuration (e.g. group.id, auto.offset.reset). + """ + GCN_KAFKA_CLIENT_ID: str = Field(min_length=1) # don't accept an empty string + GCN_KAFKA_CLIENT_SECRET: str = Field(min_length=1) + DOMAIN: str = 'gcn.nasa.gov' + KAFKA_CONFIG: dict = {} # formerly OPTIONS['CONFIG']; passed to Consumer + + +class GCNKafkaAlertStream(AlertStream): + """AlertStream implementation for GCN Kafka. + + GCN (General Coordinates Network, https://gcn.nasa.gov) distributes transient + alerts from gamma-ray and gravitational-wave observatories. This implementation + uses the gcn-kafka Python client. + + Configuration example (settings.py ALERT_STREAMS entry): + { + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.gcn.GCNKafkaAlertStream', + 'OPTIONS': { + 'GCN_KAFKA_CLIENT_ID': os.environ.get('GCN_KAFKA_CLIENT_ID', ''), + 'GCN_KAFKA_CLIENT_SECRET': os.environ.get('GCN_KAFKA_CLIENT_SECRET', ''), + 'DOMAIN': 'gcn.nasa.gov', # optional + 'KAFKA_CONFIG': {}, # optional Confluent Kafka settings + 'TOPIC_HANDLERS': { + 'gcn.circulars': 'tom_alertstreams.alertstreams.gcn.alert_logger', + 'gcn.classic.text.LVC_INITIAL': 'tom_alertstreams.alertstreams.gcn.alert_logger', + }, + }, + } + + See https://gcn.nasa.gov/docs/client for gcn_kafka client details. """ - # Upon __init__, the AlertStream base class creates instance properties from - # the settings OPTIONS dictionary, converting the keys to lowercase. - required_keys = ['GCN_CLASSIC_CLIENT_ID', 'GCN_CLASSIC_CLIENT_SECRET', 'TOPIC_HANDLERS'] - allowed_keys = ['GCN_CLASSIC_CLIENT_ID', 'GCN_CLASSIC_CLIENT_SECRET', 'TOPIC_HANDLERS', 'DOMAIN', 'CONFIG'] + configuration_class = GCNConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'gcn' + + def normalize_alert(self, raw_alert: Any, topic: str = '') -> NormalizedAlert: + """Extract common fields from a GCN Kafka message (cimpl.Message). + + GCN alert payloads vary by topic — they can be VOEvent XML, plain text, or JSON. + This implementation attempts JSON decoding; if the payload is not JSON, the raw + bytes are stored as a string in raw_payload for downstream handlers to interpret. + + Args: + raw_alert: A confluent_kafka.cimpl.Message object. + topic: The GCN Kafka topic (e.g. 'gcn.classic.text.LVC_INITIAL'). + + Returns: + NormalizedAlert with stream_name, topic, and raw_payload populated. + Astronomical coordinates are not available from GCN Classic text format. + """ + # Try to decode the alert payload. GCN topics vary in format: some are JSON, + # others are plain text (VOEvent, GCN Circular text, etc.). + + # TODO: use json-schema to validate alert!! + raw_payload: dict = {} + value_bytes = raw_alert.value() + if value_bytes: + try: + raw_payload = json.loads(value_bytes.decode('utf-8')) + except (json.JSONDecodeError, UnicodeDecodeError): + raw_payload = {'raw_text': value_bytes.decode('utf-8', errors='replace')} + + # alert_id: GCN Circulars carry their circular number in 'circularId' — this is the + # value GCNPresenter uses to build the gcn.nasa.gov/circulars/{id} link. Other GCN + # topics fall back to the Kafka message key, then to 'alert_datetime'. The latter is + # the only useful field a gcn.heartbeat carries (it has just $schema + alert_datetime, + # no id); using it avoids the alternative of a process-seeded hash of the payload, + # which changes every run and is meaningless noise in the table. + circular_id = raw_payload.get('circularId') if isinstance(raw_payload, dict) else None + key_bytes = raw_alert.key() + if circular_id is not None: + alert_id = str(circular_id) + elif key_bytes: + alert_id = key_bytes.decode('utf-8') + else: + alert_id = raw_payload.get('alert_datetime', '') if isinstance(raw_payload, dict) else '' + + # GCN alerts are reports/notices, not single observations, so observation_time is + # None. published_time is when GCN issued the alert: Circulars use 'createdOn' + # (Unix epoch ms), GCN Notices v4+ and heartbeats use 'alert_datetime' (ISO 8601). + # Fall back to now() for formats with neither (VOEvent XML, plain text). + published_time = datetime.now(timezone.utc) + if isinstance(raw_payload, dict): + if raw_payload.get('createdOn') is not None: + published_time = datetime.fromtimestamp(raw_payload['createdOn'] / 1000, tz=timezone.utc) + elif raw_payload.get('alert_datetime'): + published_time = datetime.fromisoformat(raw_payload['alert_datetime']) + + normalized_alert = NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic or raw_alert.topic(), + observation_time=None, + published_time=published_time, + alert_id=alert_id, + raw_payload=raw_payload, + ) + return normalized_alert + + def listen(self) -> None: + """Consume GCN Kafka alerts and dispatch to configured topic handlers. + + Runs an infinite loop consuming messages from the GCN Kafka broker. Kafka + errors are logged but do not terminate the loop — transient connectivity + issues should self-heal on the next consume() call. + """ + # configure and instanciate the Consumer + consumer = Consumer( + client_id=self.config.GCN_KAFKA_CLIENT_ID, + client_secret=self.config.GCN_KAFKA_CLIENT_SECRET, + domain=self.config.DOMAIN, + config=self.config.KAFKA_CONFIG, + ) + # subscribe to the topics specified in the configuration + consumer.subscribe(list(self.config.TOPIC_HANDLERS.keys())) + + while True: + for alert in consumer.consume(): + kafka_error = alert.error() + if kafka_error is not None: + logger.error( + f'GCNClassicAlertStream KafkaError: {kafka_error.name()}: {kafka_error.str()}' + ) + continue + + topic = alert.topic() + if topic not in self.alert_handler: + # this shouldn't happen be cause we subscribe to the topics for which + # we configured alert handlers for, but just in case: + logger.error( + f'GCNClassicAlertStream: alert from topic "{topic}" received ' + f'but no handler defined. Configured topics: {list(self.alert_handler.keys())}' + ) + continue + + # Unified handler convention: alert_stream=self for DI, topic=topic + # so the handler can pass it to normalize_alert(). + self.alert_handler[topic](alert, alert_stream=self, topic=topic) - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - # properties have been created from the OPTIONS dicttionary + consumer.close() - def listen(self): - super().listen() - consumer = Consumer(client_id=self.gcn_classic_client_id, - client_secret=self.gcn_classic_client_secret, - domain=self.domain, - config=self.config, - ) +def alert_logger(raw_alert: Any, alert_stream: AlertStream, topic: str, **kwargs: Any) -> None: + """Example alert handler for GCNClassicAlertStream. - consumer.subscribe(list(self.topic_handlers.keys())) + Logs the topic and raw value of the Kafka message. Use this as a starting + point for writing custom handlers; copy it into your TOM's custom_code app + and modify as needed. - # logger.debug(f'Here is a list of the available topics for {self.domain}') - # for topic in consumer.list_topics().topics: - # logger.debug(f'topic: {topic}') + The **kwargs signature absorbs alert_stream, topic, and any other extras + passed by the unified handler calling convention. - # what is a cimpl.Message?, cimpl.KafkaError? - # see https://docs.confluent.io/4.1.1/clients/confluent-kafka-python/index.html#message - while True: - for alert in consumer.consume(): - kafka_error = alert.error() # cimpl.KafkaError - if kafka_error is None: - # no error, so call the alert handler - topic = alert.topic() - try: - self.alert_handler[topic](alert) - except KeyError as err: - logger.error(f'alert from topic {topic} received but no handler defined. err: {err}') - else: - logger.error(f'GCNClassicAlertStream KafkaError: {kafka_error.name()}: {kafka_error.str()}') - consumer.close() + Args: + raw_alert: A confluent_kafka.cimpl.Message object. + **kwargs: Absorbs alert_stream, topic, and stream-specific extras. + """ + logger.info(f'gcn.alert_logger topic: {raw_alert.topic()}') + logger.info(f'gcn.alert_logger value: {raw_alert.value()}') -def alert_logger(alert): - """Example alert handler for GCN Classic over Kafka +def save_heartbeat_hourly(raw_alert: Any, alert_stream: AlertStream, **kwargs: Any) -> Any: + """Alert handler that saves only the gcn.heartbeat on the hour, dropping the rest. - This alert handler simply logs the topic and value of the cimpl.Message instance. + GCN emits a heartbeat every second — far too frequent for the demo. - See https://docs.confluent.io/4.1.1/clients/confluent-kafka-python/index.html#message - for cimpl.Message details. - """ - logger.info(f'gcn.alert_logger alert.topic(): {alert.topic()}') - logger.info(f'gcn.alert_logger alert.value(): {alert.value()}') + Args: + raw_alert: The GCN Kafka message (confluent_kafka.cimpl.Message). + alert_stream: The AlertStream instance (injected; forwarded to save_alert_to_database). + **kwargs: Stream-specific extras (e.g. topic); forwarded unchanged. + Returns: + The saved Alert, or None if this heartbeat was dropped (or the save failed). + """ + # raw_alert.timestamp() -> (timestamp_type, milliseconds_since_epoch). is_in_hourly_window + # keeps just the one heartbeat in each UTC hour's first second (see its docstring). + _, timestamp_ms = raw_alert.timestamp() + if is_in_hourly_window(timestamp_ms): + return save_alert_to_database(raw_alert, alert_stream=alert_stream, **kwargs) + return None # not in the hour's first second — drop this heartbeat diff --git a/tom_alertstreams/alertstreams/hopskotch.py b/tom_alertstreams/alertstreams/hopskotch.py index 3ec9518..26d7003 100644 --- a/tom_alertstreams/alertstreams/hopskotch.py +++ b/tom_alertstreams/alertstreams/hopskotch.py @@ -1,178 +1,302 @@ -from datetime import datetime, timezone +from __future__ import annotations + import logging import re -import uuid import traceback +import uuid +from datetime import datetime, timezone +from typing import Any, ClassVar -from django.utils import timezone as tz from django.core.exceptions import ImproperlyConfigured +from django.utils import timezone as tz from hop import Stream from hop.auth import Auth -from hop.models import JSONBlob from hop.io import Metadata, StartPosition, list_topics +from hop.models import JSONBlob -from tom_alertstreams.alertstreams.alertstream import AlertStream +from tom_alertstreams.alertstreams.alertstream import AlertStream, AlertStreamConfig, NormalizedAlert logger = logging.getLogger(__name__) -class HopskotchAlertStream(AlertStream): +class HopskotchConfig(AlertStreamConfig): + """Pydantic configuration model for HopskotchAlertStream. + + Inherits from AlertStreamConfig (a Pydantic BaseModel), so Pydantic validates + that URL, GROUP_ID, USERNAME, and PASSWORD are present. + + Fields: + URL: Hopskotch broker URL (required). Typically 'kafka://kafka.scimma.org/'. + GROUP_ID: Kafka consumer group ID (required). Must be prefixed with your + SCiMMA username to match SCiMMA Auth permissions. Format: + '-'. + USERNAME: SCiMMA Auth username (required). Obtain at https://hop.scimma.org/. + PASSWORD: SCiMMA Auth password (required). Obtain at https://hop.scimma.org/. + TOPIC_HANDLERS: Inherited from AlertStreamConfig. Maps Hopskotch topic names + to handler dotted-paths. Supports wildcards: '*' matches all public topics; + 'prefix.*' matches topics whose names match the regex. + START_POSITION: Where to start consuming. 'LATEST' (default) means only new + messages; 'EARLIEST' replays from the beginning of the topic's retention window. """ + URL: str + GROUP_ID: str + USERNAME: str + PASSWORD: str + START_POSITION: str = 'LATEST' + + +class HopskotchAlertStream(AlertStream): + """AlertStream implementation for SCiMMA Hopskotch (hop.scimma.org). + + Hopskotch is a Kafka-based message bus for time-domain astronomy operated by + SCiMMA (https://scimma.org). It carries alerts from multiple sources including + HERMES and GW notices. This implementation uses the hop-client Python library. + + Special topic support: + - '*' in TOPIC_HANDLERS subscribes to ALL public topics via the wildcard handler. + - 'prefix.*' patterns subscribe to all public topics matching the regex. + - Direct topic names take priority over wildcard matches. + + Configuration example (settings.py ALERT_STREAMS entry): + { + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.hopskotch.HopskotchAlertStream', + 'OPTIONS': { + 'URL': 'kafka://kafka.scimma.org/', + 'GROUP_ID': os.environ.get('SCIMMA_AUTH_USERNAME', '') + '-my-tom', + 'USERNAME': os.environ.get('SCIMMA_AUTH_USERNAME', ''), + 'PASSWORD': os.environ.get('SCIMMA_AUTH_PASSWORD', ''), + 'START_POSITION': 'LATEST', # optional + 'TOPIC_HANDLERS': { + 'sys.heartbeat': 'tom_alertstreams.alertstreams.hopskotch.heartbeat_handler', + 'hermes.*': 'tom_alertstreams.alertstreams.hopskotch.alert_logger', + }, + }, + } + + See https://hop-client.readthedocs.io/ for hop-client documentation. """ - required_keys = ['URL', 'GROUP_ID', 'USERNAME', 'PASSWORD', 'TOPIC_HANDLERS'] - allowed_keys = ['URL', 'GROUP_ID', 'USERNAME', 'PASSWORD', 'TOPIC_HANDLERS', 'START_POSITION'] - PUBLIC_TOPIC_CHECK_INTERVAL = 300 # Seconds between checking for new public topics - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - # the following methods may fail if improperly configured. - # So, do them now to catch any errors, before listen() is spawned in it's own Process. - logger.debug(f'HopskotchAlertStream.__init__() kwargs: {kwargs}') + configuration_class = HopskotchConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'hopskotch' + + # Seconds between checks for new public topics when wildcard subscriptions are active. + PUBLIC_TOPIC_CHECK_INTERVAL: ClassVar[int] = 300 + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + logger.debug(f'HopskotchAlertStream.__init__() config: {self.config}') + + # Fetch public topics and build the stream URL up front — if the configuration + # is broken, we want to fail here (before listen() spawns in its own thread) + # so the error is visible immediately at startup. self.public_topics = self.get_all_public_topics() self.stream_url = self.get_stream_url() start_position = StartPosition.LATEST - if hasattr(self, 'start_position') and self.start_position == 'EARLIEST': + if self.config.START_POSITION == 'EARLIEST': start_position = StartPosition.EARLIEST self.stream = self.get_stream(start_position) def get_all_public_topics(self) -> list[str]: - """Returns the up-to-date list of Topic names to consume. + """Return the current list of publicly-readable Hopskotch topic names. - Use the saved options to repeatedly construct the topic list, and - keep it in sync with the publicaly_readable topics from SCiMMA Auth. + Queries the Hopskotch broker directly via the hop-client. Filters out + internal Kafka topics (those starting with '__' or containing no '.'). - The Topic list is a combination of the - a. the publicly_readable Topics from SCiMMA Auth - b. any topics supplied on the command line via -T, --topic + Returns: + List of topic name strings available on the Hopskotch broker. """ - hop_auth = Auth(self.username, self.password) - logger.info('getting publicly_readable topics from SCiMMA Auth.') - # use the hop-client to ask Kafka directly for the topics since SCiMMA Auth can be out of sync - # include only topics that a) contain a '.'; b) don't start with '__' (excludes __consumer_offsets) - publicly_readable_topics = [topic for topic in list_topics(self.url, hop_auth).keys() - if not (topic.startswith('__') and (topic.count('.')==0))] - logger.debug(f'publicly_readable_topics: {publicly_readable_topics}') - - return publicly_readable_topics + hop_auth = Auth(self.config.USERNAME, self.config.PASSWORD) + logger.info('HopskotchAlertStream: fetching public topics from SCiMMA Auth.') + all_topics = list_topics(self.config.URL, hop_auth) + # Exclude internal Kafka topics (__consumer_offsets, etc.) and topics + # without a namespace separator (no '.') which are internal by convention. + publicly_readable = [ + topic for topic in all_topics.keys() + if not (topic.startswith('__') and topic.count('.') == 0) + ] + logger.debug(f'HopskotchAlertStream public topics: {publicly_readable}') + return publicly_readable def get_stream_url(self) -> str: - """For Hopskotch, topics are specified on the url. So, this - method gets a base url (from super) and then adds topics to it. + """Build the Hopskotch stream URL with topics appended. - Hopskotch (hop.io) requires at least one topic to be specified. + Hopskotch requires topics to be specified in the URL rather than via a + subscribe() call. This method resolves wildcard patterns against the + current public topic list before building the URL. - You might not need a method like this if your Kafka client provides - alternative ways to subscribe to a topic. For example, the gcn_kafka.Consumer - class provides a 'substribe([list of topics])' method. (see gcn.py). + Returns: + Fully-qualified Hopskotch stream URL with topics. + + Raises: + ImproperlyConfigured: if TOPIC_HANDLERS is empty (hop requires ≥1 topic). """ - logger.debug(f'HopskotchAlertStream.get_stream_url topics: {list(self.topic_handlers.keys())}') - if self.topic_handlers == {}: - msg = 'Hopskotch requires at least one topic to open the stream. Check ALERT_STREAMS in settings.py' - raise ImproperlyConfigured(msg) - - base_stream_url = self.url - - # if not present, add trailing slash to base_stream url - # so, comma-separated topics can be appeneded. - if base_stream_url[-1] != '/': - base_stream_url += '/' - - # append comma-separated topics to base URL - specified_topics = set(self.topic_handlers.keys()) - if '*' in specified_topics: - # Add all public topics if a asterisk is set in the topic_handlers - specified_topics = specified_topics.union(set(self.public_topics)) + if not self.config.TOPIC_HANDLERS: + raise ImproperlyConfigured( + 'HopskotchAlertStream requires at least one entry in TOPIC_HANDLERS. ' + 'Check ALERT_STREAMS in settings.py.' + ) + + base_url = self.config.URL.rstrip('/') + '/' + + # Expand wildcard patterns against the public topic list. + specified = set(self.config.TOPIC_HANDLERS.keys()) + if '*' in specified: + # Full wildcard: subscribe to every public topic. + specified = specified | set(self.public_topics) else: - # Look over all topics, and if there are any with a partial wildcard in them, - # Add all the public topics that match that partial wildcard - for topic in specified_topics: - if '*' in topic: - specified_topics = specified_topics.union( - set([t for t in self.public_topics if re.match(topic, t)])) - - # Also remove topics with wildcards in them, and convert specified topics set to list - specified_topics = [topic for topic in specified_topics if not '*' in topic] - - topics = ','.join(specified_topics) # 'topic1,topic2,topic3' - hopskotch_stream_url = base_stream_url + topics - - logger.debug(f'HopskotchAlertStream.get_stream_url url: {hopskotch_stream_url}') - return hopskotch_stream_url - - def get_stream(self, start_position=StartPosition.LATEST) -> Stream: - hop_auth = Auth(self.username, self.password) - - # TODO: allow StartPosition to be set from OPTIONS configuration dictionary - stream = Stream(auth=hop_auth, start_at=start_position) - return stream - - def listen(self): - super().listen() - # TODO: alternatively, WARN upon OPTIONS['topics'] extries that don't have - # handlers in the alert_handler. (i.e they've configured a topic subscription - # without providing a handler for the topic. So, warn them). + # Partial wildcards: 'hermes.*' → all public topics matching the regex. + for pattern in list(specified): + if '*' in pattern: + specified |= {t for t in self.public_topics if re.match(pattern, t)} + + # Remove the wildcard placeholders — real topic names only. + concrete_topics = [t for t in specified if '*' not in t] + hopskotch_url = base_url + ','.join(concrete_topics) + logger.debug(f'HopskotchAlertStream stream URL: {hopskotch_url}') + return hopskotch_url + + def get_stream(self, start_position: StartPosition = StartPosition.LATEST) -> Stream: + """Create and return a hop-client Stream object. + + Args: + start_position: Where to start consuming (LATEST or EARLIEST). + + Returns: + An authenticated hop.Stream ready for use in listen(). + """ + hop_auth = Auth(self.config.USERNAME, self.config.PASSWORD) + return Stream(auth=hop_auth, start_at=start_position) + + def normalize_alert(self, raw_alert: Any, topic: str = '') -> NormalizedAlert: + """Extract common fields from a Hopskotch JSONBlob alert. + + Hopskotch delivers alerts as hop.models.JSONBlob objects (or other hop model + types). The .content attribute holds the parsed dict. Since Hopskotch carries + alerts from many sources (HERMES, GW notices, etc.), only a generic extraction + is possible at this level; science-specific handlers should subclass and override. + + Args: + raw_alert: A hop.models.JSONBlob (or similar hop model) object. + topic: The Hopskotch topic the alert arrived on. + + Returns: + NormalizedAlert with stream_name, topic, and raw_payload populated. + """ + content = getattr(raw_alert, 'content', None) or {} + alert_id = str(content.get('message_id', id(raw_alert))) + return NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic, + observation_time=None, # generic Hopskotch messages carry no parsed obs/publish time + published_time=None, + alert_id=alert_id, + raw_payload=content if isinstance(content, dict) else {}, + ) + + def listen(self) -> None: + """Consume Hopskotch alerts and dispatch to configured topic handlers. + + Runs an infinite loop reading from the Hopskotch stream. Periodically checks + for new public topics (every PUBLIC_TOPIC_CHECK_INTERVAL seconds) and restarts + the stream if the topic list has changed, so new topics are picked up without + a manual restart. + + Topic matching priority: + 1. Exact topic name match + 2. Regex wildcard pattern match (e.g. 'hermes.*') + 3. Catch-all '*' handler + + Handler calling convention: + handler(alert, alert_stream=self, topic=topic, metadata=metadata) + Handlers absorb extras they do not need via **kwargs. + """ last_check_time = tz.now() while True: try: - logger.info(f'HopskotchAlertStream.listen opening stream: {self.stream_url} with group_id: {self.group_id}') - with self.stream.open(self.stream_url, 'r', group_id=self.group_id) as src: + logger.info( + f'HopskotchAlertStream: opening stream {self.stream_url} ' + f'with group_id: {self.config.GROUP_ID}' + ) + with self.stream.open(self.stream_url, 'r', group_id=self.config.GROUP_ID) as src: for alert, metadata in src.read(metadata=True): - # type(gcn_circular) is - # type(metadata) is - if metadata.topic in self.alert_handler: - # TODO: should probably use *args, **kwargs to pass unknow number of arguments - self.alert_handler[metadata.topic](alert, metadata) + topic = metadata.topic + + # Determine the handler: exact match, then regex wildcard, then '*'. + if topic in self.alert_handler: + handler = self.alert_handler[topic] else: - # First check all wildcard topics to see if they will match this topic - matched_handler = False - for topic in self.alert_handler.keys(): - if topic != '*' and '*' in topic and re.match(topic, metadata.topic): - self.alert_handler[topic](alert, metadata) - matched_handler = True + handler = None + for pattern, candidate in self.alert_handler.items(): + if pattern != '*' and '*' in pattern and re.match(pattern, topic): + handler = candidate break - if not matched_handler: - # If nothing matched and we have a catch all handler, fall back to default public topic handler - if '*' in self.alert_handler: - self.alert_handler['*'](alert, metadata) - else: - # TODO: should define a default handler for all unhandeled topics - logger.error(f'alert from topic {metadata.topic} received but no handler defined.') + if handler is None: + handler = self.alert_handler.get('*') + + if handler is not None: + # Unified convention + Hopskotch-specific metadata kwarg. + handler(alert, alert_stream=self, topic=topic, metadata=metadata) + else: + logger.error( + f'HopskotchAlertStream: alert from topic "{topic}" received ' + f'but no handler matched. Configured: {list(self.alert_handler.keys())}' + ) + + # Periodically refresh public topics to pick up new ones automatically. if (tz.now() - last_check_time).total_seconds() > self.PUBLIC_TOPIC_CHECK_INTERVAL: last_check_time = tz.now() - public_topics = self.get_all_public_topics() - if set(public_topics) != set(self.public_topics): - logger.info(f"New public topics found, restarting hop stream") - self.public_topics = public_topics + fresh_topics = self.get_all_public_topics() + if set(fresh_topics) != set(self.public_topics): + logger.info('HopskotchAlertStream: new public topics found — restarting stream.') + self.public_topics = fresh_topics self.stream_url = self.get_stream_url() - break + break # Exit inner loop; outer while True reopens the stream. + except Exception as ex: logger.error(f'HopskotchAlertStream.listen: {ex}') - logger.error(traceback.format_exc()) # Show the traceback so we have a chance of figuring out what is breaking + logger.error(traceback.format_exc()) -def heartbeat_handler(heartbeat: JSONBlob, metadata: Metadata): - """Example alert handler for HopskotchAlertStream sys.heartbeat topic. - Note that HopskotchAlertStream.listen() method knows that Hopskotch alerts come with - both alert and metadata. So, the alert_handler methods have a signiture (taking both - as arguments) specific to this stream. +def heartbeat_handler(heartbeat: JSONBlob, **kwargs: Any) -> None: + """Example handler for the Hopskotch sys.heartbeat topic. + + Logs every 300th heartbeat to avoid flooding the log. Copy into your TOM's + custom_code app and modify as needed. + + The **kwargs signature absorbs alert_stream, topic, metadata, and any other + extras passed by the unified handler calling convention. + + Args: + heartbeat: A hop.models.JSONBlob with a 'timestamp' and 'count' in .content. + **kwargs: Absorbs alert_stream, topic, metadata, and stream-specific extras. """ - content: dict = heartbeat.content # see hop_client reatthedocs - timestamp = datetime.fromtimestamp(content["timestamp"] / 1e6, tz=timezone.utc) - if heartbeat.content['count'] % 300 == 0: - # mod 300 just for convenience so as not to flood logger - logging.info(f'{timestamp.isoformat()} heartbeat.content dict: {heartbeat.content}. metadata: {metadata}') + content: dict = heartbeat.content + timestamp = datetime.fromtimestamp(content['timestamp'] / 1e6, tz=timezone.utc) + if content.get('count', 0) % 300 == 0: + logger.info(f'Hopskotch heartbeat at {timestamp.isoformat()}: {content}') + + +def alert_logger(alert: JSONBlob, **kwargs: Any) -> None: + """Example alert handler for HopskotchAlertStream. + + Logs the topic and alert UUID. Copy into your TOM's custom_code app and + modify as needed. + The **kwargs signature absorbs alert_stream, topic, metadata, and any other + extras passed by the unified handler calling convention. Access metadata via + kwargs.get('metadata') if needed. -def alert_logger(alert: JSONBlob, metadata: Metadata): - """Example alert handler. The method signsture is specific to Hopskotch alerts. + Args: + alert: A hop.models.JSONBlob (or other hop model type). + **kwargs: Absorbs alert_stream, topic, metadata, and stream-specific extras. """ - # search the header (list of tuples) for a UUID-tuple (keyed by '_id') - # eg. ('_id', b'$\xd6oGmVM\xed\x97\xe7|\x1c\x8f\x11V\xe9') - alert_uuid_tuple = next((item for item in metadata.headers if item[0] == '_id'), None) - if alert_uuid_tuple: - alert_uuid = uuid.UUID(bytes=alert_uuid_tuple[1]) - else: - # in this case the alert was probably published with hop-client<0.8.0 - alert_uuid = None - logger.info(f'Alert (uuid={alert_uuid}) received on topic {metadata.topic}: {alert}; metatdata: {metadata}') + metadata: Metadata | None = kwargs.get('metadata') + alert_uuid = None + if metadata is not None: + uuid_tuple = next((h for h in metadata.headers if h[0] == '_id'), None) + if uuid_tuple: + alert_uuid = uuid.UUID(bytes=uuid_tuple[1]) + topic = kwargs.get('topic', getattr(metadata, 'topic', 'unknown') if metadata else 'unknown') + logger.info(f'Hopskotch alert (uuid={alert_uuid}) on topic "{topic}": {alert}') \ No newline at end of file diff --git a/tom_alertstreams/alertstreams/lasair.py b/tom_alertstreams/alertstreams/lasair.py new file mode 100644 index 0000000..301f17c --- /dev/null +++ b/tom_alertstreams/alertstreams/lasair.py @@ -0,0 +1,263 @@ +from __future__ import annotations + +import json +import logging +import random +import time +from datetime import datetime, timezone +from typing import Any, ClassVar + +from lasair import lasair_consumer + +from tom_alertstreams.alertstreams.alertstream import ( + AlertStream, AlertStreamConfig, NormalizedAlert, _mjd_to_datetime, +) + +logger = logging.getLogger(__name__) + +# Mock data constants — chosen to be unmistakably non-astronomical: +# (0, 0) is not a real survey pointing; 99.0 is the astronomical sentinel for "no data". +_MOCK_RA = 0.0 +_MOCK_DEC = 0.0 +_MOCK_MAGNITUDE = 99.0 +INTER_ALERT_SLEEP_MIN = 360 # six minutes +INTER_ALERT_SLEEP_MAX = 420 # seven minutes + + +# --------------------------------------------------------------------------- +# Pydantic configuration models +# --------------------------------------------------------------------------- + +class LasairMockConfig(AlertStreamConfig): + """Pydantic configuration model for LasairMockAlertStream. + + Inherits TOPIC_HANDLERS from AlertStreamConfig (a Pydantic BaseModel). + No additional fields needed for mock data generation. + """ + pass + + +class LasairConfig(AlertStreamConfig): + """Pydantic configuration for LasairAlertStream. + + Fields: + LASAIR_TOKEN: REST API token (optional — not used for Kafka auth, but + available for future REST API features like topic discovery). + Obtained from lasair.lsst.ac.uk/profile. + LASAIR_KAFKA_SERVER: Kafka broker address (required). + LSST: 'lasair-lsst-kafka.lsst.ac.uk:9092' + LASAIR_GROUP_ID: Kafka consumer group ID (required). Same group ID + resumes from the last offset; a new group ID fetches ~7 days of + cached alerts. + TOPIC_HANDLERS: Inherited from AlertStreamConfig. Maps a single Lasair + streaming filter topic to a handler dotted-path. Topics are + user-created filters from https://lasair.lsst.ac.uk/filters/. + """ + LASAIR_TOKEN: str | None = None + LASAIR_KAFKA_SERVER: str + LASAIR_GROUP_ID: str + + +# --------------------------------------------------------------------------- +# Mock Lasair stream (for demo use without Kafka connectivity) +# --------------------------------------------------------------------------- + +class LasairMockAlertStream(AlertStream): + """Mock Lasair AlertStream that generates obviously-fake alerts. + + Provides a demo fallback when Kafka connectivity is not available. + Uses the same mock pattern as the other stub streams (sentinel coordinates, + sentinel magnitude, 6–7 minute sleep between alerts). + + The mock stream uses STREAM_NAME='lasair' so it occupies the same dashboard + slot as the real Lasair stream would if it were the only one configured. + """ + configuration_class = LasairMockConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'lasair' + IS_MOCK: ClassVar[bool] = True + + def normalize_alert(self, raw_alert: dict, topic: str = '') -> NormalizedAlert: + """Map a mock Lasair alert dict to a NormalizedAlert. + + Args: + raw_alert: Dict produced by listen(); contains mock field values. + topic: Topic the alert was generated for. + + Returns: + NormalizedAlert populated from the mock dict fields. + """ + normalized_alert = NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic or raw_alert.get('topic', ''), + observation_time=None, # a mock alert has no real observation + published_time=datetime.fromisoformat(raw_alert['timestamp']), + alert_id=raw_alert['alert_id'], + object_id=raw_alert.get('object_id'), + ra=raw_alert.get('ra'), + dec=raw_alert.get('dec'), + magnitude=raw_alert.get('magnitude'), + raw_payload=raw_alert, + ) + return normalized_alert + + def listen(self) -> None: + """Generate mock Lasair alerts and dispatch to configured topic handlers. + + Loops indefinitely, emitting one mock alert per iteration with a random + 6–7 minute delay. Topics are round-robined if multiple are configured. + """ + counter = 0 + topics = list(self.config.TOPIC_HANDLERS.keys()) + + while True: + counter += 1 + topic = topics[counter % len(topics)] + timestamp = datetime.now(timezone.utc) + object_id = f'MOCK-{self.STREAM_NAME.upper()}-{counter:04d}' + alert_id = f'MOCK-{timestamp.strftime("%Y%m%d%H%M%S")}' + mock_alert: dict[str, Any] = { + 'alert_id': alert_id, + 'object_id': object_id, + 'topic': topic, + 'timestamp': timestamp.isoformat(), + 'ra': _MOCK_RA, + 'dec': _MOCK_DEC, + 'magnitude': _MOCK_MAGNITUDE, + 'mock': True, + } + logger.debug(f'LasairMockAlertStream: mock alert {object_id}') + self.alert_handler[topic](mock_alert, alert_stream=self, topic=topic) + time.sleep(random.uniform(INTER_ALERT_SLEEP_MIN, INTER_ALERT_SLEEP_MAX)) + + +# --------------------------------------------------------------------------- +# Real Lasair LSST stream +# --------------------------------------------------------------------------- + +class LasairAlertStream(AlertStream): + """Lasair LSST alert stream via Kafka. + + Uses the ``lasair`` library's ``lasair_consumer`` class to receive JSON-encoded + alerts from the Lasair LSST Kafka broker. Kafka consuming is unauthenticated + (public), unlike the REST API which requires a token. + + ``lasair_consumer`` accepts a single topic per instance, following the documented + pattern at lasair.readthedocs.io/en/main/core_functions/alert-streams.html. + Configure one TOPIC_HANDLERS entry corresponding to a streaming filter you've + created at lasair.lsst.ac.uk. + + Lasair alert JSON comes in three tiers (configured per-filter in the Lasair web UI): + - Standard: ``{diaObjectId, ra, decl, UTC}`` + - Lite: adds ``alert.diaSourcesList[]`` with psfFlux, midpointMjdTai, band, etc. + - Full: adds ``diaObject`` (80+ attrs) and comprehensive ``diaSource`` data. + + ``normalize_alert()`` handles all three tiers defensively, extracting what's available. + + Configuration example (settings.py ALERT_STREAMS entry):: + + { + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.lasair.LasairAlertStream', + 'OPTIONS': { + 'LASAIR_TOKEN': os.environ.get('LASAIR_LSST_TOKEN', ''), + 'LASAIR_KAFKA_SERVER': os.environ.get('LASAIR_KAFKA_SERVER', ''), + 'LASAIR_GROUP_ID': os.environ.get('LASAIR_GROUP_ID', ''), + 'TOPIC_HANDLERS': { + 'lasair_114_TOMToolkit.tom-demo': 'tom_alertstreams.alertstreams.alertstream.save_alert_to_database', + }, + }, + } + """ + configuration_class = LasairConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'lasair' + + def listen(self) -> None: + """Consume Lasair alerts from Kafka and dispatch to the configured handler. + + Creates a ``lasair_consumer`` for the single configured topic. The consumer + wraps ``confluent_kafka.Consumer`` with Lasair's standard settings + (``auto.offset.reset: smallest``). Polls indefinitely; the consumer is closed + on errors, KeyboardInterrupt, or normal exit via the finally block. + """ + # lasair_consumer accepts a single topic — take the one configured topic + topic = list(self.config.TOPIC_HANDLERS.keys())[0] + logger.info(f'{self.STREAM_NAME}: connecting to {self.config.LASAIR_KAFKA_SERVER}, ' + f'group_id={self.config.LASAIR_GROUP_ID}, topic={topic}') + + lasair_kafka_consumer = lasair_consumer( + self.config.LASAIR_KAFKA_SERVER, + self.config.LASAIR_GROUP_ID, + topic, + ) + + try: + while True: + msg = lasair_kafka_consumer.poll(timeout=20) + if msg is None: + continue # timeout with no message, retry + if msg.error(): + logger.warning(f'{self.STREAM_NAME} Kafka error on {topic}: {msg.error()}') + continue + alert = json.loads(msg.value()) + logger.info(f'{self.STREAM_NAME} received alert on {topic}') + self.alert_handler[topic](alert, alert_stream=self, topic=topic) + finally: + lasair_kafka_consumer.close() + + def normalize_alert(self, raw_alert: dict, topic: str = '') -> NormalizedAlert: + """Extract common fields from a Lasair LSST alert dict. + + Handles all three Lasair message tiers defensively: checks for lite/full + tier fields first (nested ``alert.diaSourcesList``), then falls back to + standard tier fields (top-level ``UTC`` timestamp). + + Lasair field naming differs from other LSST streams: + - ``decl`` instead of ``dec`` + - ``diaObjectId`` at top level (not nested in diaSource) + - ``UTC`` as ISO string (standard tier) vs ``midpointMjdTai`` (lite/full) + + Args: + raw_alert: Alert dict decoded from Lasair's JSON Kafka message. + topic: The Lasair topic the alert was consumed from. + + Returns: + NormalizedAlert with tier-appropriate fields populated. + """ + # diaObjectId is at top level in all Lasair tiers + dia_object_id = raw_alert.get('diaObjectId') + + # Timestamp and flux: prefer lite/full tier diaSourcesList if available, + # fall back to standard tier's UTC string + timestamp = None + flux = None + dia_sources = raw_alert.get('alert', {}).get('diaSourcesList', []) + if dia_sources: + # Lite/full tier: use the most recent diaSource entry + latest_source = dia_sources[0] + mjd = latest_source.get('midpointMjdTai') + if mjd is not None: + timestamp = _mjd_to_datetime(mjd) + flux = latest_source.get('psfFlux') + + # Fall back to standard tier UTC timestamp if lite/full didn't provide one + if timestamp is None: + utc_str = raw_alert.get('UTC') + if utc_str: + timestamp = datetime.fromisoformat(utc_str).replace(tzinfo=timezone.utc) + else: + timestamp = None # no observation time available for this alert + + normalized_alert = NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic, + observation_time=timestamp, + published_time=None, + alert_id=str(dia_object_id or ''), # Lasair has no separate alert ID + object_id=str(dia_object_id) if dia_object_id else None, + ra=raw_alert.get('ra'), + dec=raw_alert.get('decl'), # NB: Lasair uses 'decl' not 'dec' + magnitude=None, # LSST uses flux, not magnitude + flux=flux, + raw_payload=raw_alert, + ) + return normalized_alert diff --git a/tom_alertstreams/alertstreams/pittgoogle.py b/tom_alertstreams/alertstreams/pittgoogle.py new file mode 100644 index 0000000..9a31d47 --- /dev/null +++ b/tom_alertstreams/alertstreams/pittgoogle.py @@ -0,0 +1,339 @@ +from __future__ import annotations + +import logging +import random +import time +from datetime import datetime, timezone +from typing import Any, ClassVar + +import pittgoogle +from pydantic import model_validator + +from tom_alertstreams.alertstreams.alertstream import ( + AlertStream, AlertStreamConfig, NormalizedAlert, _jd_to_datetime, is_in_minute_window, + save_alert_to_database, +) + +logger = logging.getLogger(__name__) + +# Pitt-Google publishes its public topics in this Google Cloud project. +PITTGOOGLE_PROJECT_DEFAULT: str = pittgoogle.ProjectIds().pittgoogle # 'ardent-cycling-243415' + +# Mock data constants — chosen to be unmistakably non-astronomical: +# (0, 0) is not a real survey pointing; 99.0 is the astronomical sentinel for "no data". +_MOCK_RA = 0.0 +_MOCK_DEC = 0.0 +_MOCK_MAGNITUDE = 99.0 +INTER_ALERT_SLEEP_MIN = 3600 # sixty minutes +INTER_ALERT_SLEEP_MAX = 3600 + + +def _schema_for_topic(topic_name: str) -> str: + """Return the pittgoogle alert schema used to deserialize a topic's messages.""" + if topic_name.startswith('lsst'): + return 'lsst' + if topic_name.startswith('lvk'): + return 'lvk' + return 'ztf' + + +# --------------------------------------------------------------------------- +# Pydantic configuration models +# --------------------------------------------------------------------------- + +class PittGoogleMockConfig(AlertStreamConfig): + """Pydantic configuration model for PittGoogleMockAlertStream. + + Inherits TOPIC_HANDLERS from AlertStreamConfig (a Pydantic BaseModel). + No additional fields needed for mock data generation. + """ + pass + + +class PittGoogleConfig(AlertStreamConfig): + """Pydantic configuration for the real PittGoogleAlertStream. + + Pitt-Google distributes alerts over Google Cloud Pub/Sub (not Kafka). The topic + is the single key in TOPIC_HANDLERS (as with GCN/Hopskotch); a Pub/Sub subscription + binds to exactly one topic, so each entry handles one topic. GCP authentication is + read from the environment by pittgoogle — not configured here: + GOOGLE_CLOUD_PROJECT the project that owns the subscription + GOOGLE_APPLICATION_CREDENTIALS path to the service-account key JSON + + Fields: + PITTGOOGLE_PROJECT: Google Cloud project that publishes the topic. Defaults to + Pitt-Google's public project ('ardent-cycling-243415'). + PITTGOOGLE_SUBSCRIPTION: Name of the (persistent) subscription to create in our + project. Optional; defaults to the topic name, so two entries that read + different topics get distinct subscriptions automatically. + TOPIC_HANDLERS: Inherited from AlertStreamConfig. Must contain exactly one topic. + """ + PITTGOOGLE_PROJECT: str = PITTGOOGLE_PROJECT_DEFAULT + PITTGOOGLE_SUBSCRIPTION: str | None = None + + @model_validator(mode='after') + def _check_single_topic(self) -> PittGoogleConfig: + """Require exactly one topic per entry. + + A Pub/Sub subscription binds to exactly one topic and pittgoogle.Consumer.stream() + consumes a single subscription, so a PittGoogleAlertStream instance handles one + topic. To consume several Pitt-Google topics, configure several ALERT_STREAMS + entries (one per topic) — readstreams runs each in its own thread. + """ + if len(self.TOPIC_HANDLERS) != 1: + raise ValueError( + f'PittGoogleAlertStream requires exactly one topic per entry (one Pub/Sub ' + f'subscription binds to one topic); got {len(self.TOPIC_HANDLERS)}: ' + f'{list(self.TOPIC_HANDLERS)}. Configure a separate ALERT_STREAMS entry per topic.' + ) + return self + + +# --------------------------------------------------------------------------- +# Throttling alert handler for high-rate Pitt-Google topics (ztf-loop, ztf-alerts) +# --------------------------------------------------------------------------- + +def save_pittgoogle_throttled(raw_alert: Any, alert_stream: AlertStream, **kwargs: Any) -> Any: + """Alert handler that saves only ~one alert per UTC minute, dropping (but acking) the rest. + + Throttles a high-rate Pitt-Google topic to ~one saved alert per minute. Both ZTF topics + need this: ztf-loop replays a recent alert ~1/sec, and ztf-alerts is the full ZTF firehose + (tens/sec when ZTF observes, plus any accumulated subscription backlog). Saving every alert + would swamp the demo's SQLite DB and — worse — outrun Pub/Sub's lease deadline: the client + can't ack fast enough, leases expire, and Pub/Sub redelivers ("Dropping N items because they + were leased too long"). Dropping-but-acking the rest keeps acks fast, so the consumer stays + current and any backlog drains. The demo only displays the recent few per topic anyway. + Mirrors gcn.save_heartbeat_hourly; a minute cadence here keeps the live ZTF feed visibly + fresh. To change the cadence, swap is_in_minute_window for is_in_hourly_window. + + The throttle clock is the Pub/Sub publishTime (raw_alert.msg.publish_time) gated by + is_in_minute_window. (Deliberately distinct from the alert's published_time field, which + records the survey's upstream kafka.timestamp; see PittGoogleAlertStream.normalize_alert.) + + Args: + raw_alert: The pittgoogle.Alert delivered by the Consumer. + alert_stream: The AlertStream instance (injected; forwarded to save_alert_to_database). + **kwargs: Stream-specific extras (e.g. topic); forwarded unchanged. + + Returns: + The saved Alert, or None if this alert was dropped (or the save failed). The + Consumer acks the Pub/Sub message either way (see PittGoogleAlertStream.listen). + """ + # Throttle on the Pub/Sub publishTime; fall back to receipt time if it's absent. + message = getattr(raw_alert, 'msg', None) + publish_time = getattr(message, 'publish_time', None) + if publish_time is None: + publish_time = datetime.now(timezone.utc) + timestamp_ms = int(publish_time.timestamp() * 1000) + if is_in_minute_window(timestamp_ms): + return save_alert_to_database(raw_alert, alert_stream=alert_stream, **kwargs) + return None # outside the minute's first second — drop this alert + + +# --------------------------------------------------------------------------- +# Mock Pitt-Google stream (demo fallback without GCP credentials) +# --------------------------------------------------------------------------- + +class PittGoogleMockAlertStream(AlertStream): + """Mock Pitt-Google AlertStream that generates obviously-fake alerts. + + Provides a demo fallback when Google Cloud credentials are not available. Uses the + same mock pattern as the other stub streams (sentinel coordinates, sentinel magnitude). + STREAM_NAME='pittgoogle' so it occupies the same dashboard slot as the real stream. + """ + configuration_class = PittGoogleMockConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'pittgoogle' + IS_MOCK: ClassVar[bool] = True # generates fake alerts, not a real Pitt-Google feed + + def normalize_alert(self, raw_alert: dict, topic: str = '') -> NormalizedAlert: + """Map a mock Pitt-Google alert dict to a NormalizedAlert. + + Args: + raw_alert: Dict produced by listen(); contains mock field values. + topic: Topic the alert was generated for. + + Returns: + NormalizedAlert populated from the mock dict fields. + """ + return NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic or raw_alert.get('topic', ''), + observation_time=None, # a mock alert has no real observation + published_time=datetime.fromisoformat(raw_alert['timestamp']), + alert_id=raw_alert['alert_id'], + object_id=raw_alert.get('object_id'), + ra=raw_alert.get('ra'), + dec=raw_alert.get('dec'), + magnitude=raw_alert.get('magnitude'), + raw_payload=raw_alert, + ) + + def listen(self) -> None: + """Generate mock Pitt-Google alerts and dispatch to configured topic handlers. + + Loops indefinitely, emitting one mock alert per iteration. Topics are + round-robined if multiple are configured. + """ + counter = 0 + topics = list(self.config.TOPIC_HANDLERS.keys()) + + # for this mock, generate fake alerts (rather than listen to the stream) endlessly + while True: + counter += 1 + topic = topics[counter % len(topics)] + timestamp = datetime.now(timezone.utc) + object_id = f'MOCK-{self.STREAM_NAME.upper()}-{counter:04d}' + alert_id = f'MOCK-{timestamp.strftime("%Y%m%d%H%M%S")}' + mock_alert: dict[str, Any] = { + 'alert_id': alert_id, + 'object_id': object_id, + 'topic': topic, + 'timestamp': timestamp.isoformat(), + 'ra': _MOCK_RA, + 'dec': _MOCK_DEC, + 'magnitude': _MOCK_MAGNITUDE, + 'mock': True, + } + logger.debug(f'PittGoogleMockAlertStream: mock alert {object_id}') + self.alert_handler[topic](mock_alert, alert_stream=self, topic=topic) + time.sleep(random.uniform(INTER_ALERT_SLEEP_MIN, INTER_ALERT_SLEEP_MAX)) + + +# --------------------------------------------------------------------------- +# Real Pitt-Google stream (Google Cloud Pub/Sub) +# --------------------------------------------------------------------------- + +class PittGoogleAlertStream(AlertStream): + """Real Pitt-Google alert stream over Google Cloud Pub/Sub (not Kafka). + + Pitt-Google distributes alerts over Pub/Sub. We consume by creating a subscription in + our own Google Cloud project (GOOGLE_CLOUD_PROJECT), attached to one of Pitt-Google's + public topics, then streaming messages through a pittgoogle.Consumer. + + A subscription binds to exactly one topic, so each instance handles a single topic (the + one key in TOPIC_HANDLERS). Configure one ALERT_STREAMS entry per topic — e.g. the + throttled 'ztf-loop' (a ~1/sec firehose, handled by save_ztf_loop_hourly) and the live + 'ztf-alerts'. Both use STREAM_NAME='pittgoogle'; readstreams runs each in its own thread. + + Authentication is read from the environment by pittgoogle: + GOOGLE_CLOUD_PROJECT our project that owns the subscription + GOOGLE_APPLICATION_CREDENTIALS path to the service-account key JSON + + Configuration example (settings.py ALERT_STREAMS entry):: + + { + 'ACTIVE': True, + 'NAME': 'tom_alertstreams.alertstreams.pittgoogle.PittGoogleAlertStream', + 'OPTIONS': { + 'TOPIC_HANDLERS': { + 'ztf-loop': 'tom_alertstreams.alertstreams.pittgoogle.save_ztf_loop_hourly', + }, + }, + } + """ + configuration_class = PittGoogleConfig # type: ignore[assignment] + STREAM_NAME: ClassVar[str] = 'pittgoogle' + + def listen(self) -> None: + """Stream alerts from one Pitt-Google Pub/Sub topic and dispatch to its handler. + + Creates (or verifies) a persistent subscription to the topic, then opens a streaming + pull via pittgoogle.Consumer. Consumer.stream() blocks and raises on a fatal error, + so the AlertStream.run() supervisor handles reconnect — no retry loop here. + """ + topic_name = next(iter(self.config.TOPIC_HANDLERS)) # one topic per entry (validated) + subscription_name = self.config.PITTGOOGLE_SUBSCRIPTION or topic_name + schema_name = _schema_for_topic(topic_name) + + subscription = pittgoogle.Subscription( + subscription_name, + topic=pittgoogle.Topic(topic_name, projectid=self.config.PITTGOOGLE_PROJECT), + schema_name=schema_name, + ) + subscription.touch() # create-or-verify in our project; persistent (NOT deleted) + logger.info(f'{self.STREAM_NAME}: streaming {topic_name} via subscription {subscription_name}') + + def message_callback(alert: pittgoogle.Alert) -> pittgoogle.pubsub.Response: + # The Consumer builds Alert(msg) WITHOUT a schema_name; set it from the topic so the + # schema-aware accessors (sourceid/objectid/dict) resolve in normalize_alert. + alert.schema_name = schema_name + # Dispatch to the configured handler (save_alert_to_database, or the ztf-loop + # throttle). ALWAYS ack afterward: Pub/Sub redelivers un-acked messages, so a + # throttled-and-dropped alert must be acked too, or it piles up as an un-acked + # backlog. We swallow handler errors here so a single bad alert is acked + logged + # rather than redelivered forever (save_alert_to_database already swallows its own). + try: + self.alert_handler[topic_name](alert, alert_stream=self, topic=topic_name) + except Exception: + logger.exception(f'{self.STREAM_NAME}: handler error on topic {topic_name}') + return pittgoogle.pubsub.Response(ack=True) + + consumer = pittgoogle.Consumer(subscription, msg_callback=message_callback) + # Open the streaming pull in the background, then block on its future ourselves. + # pittgoogle's Consumer.stream(block=True) would instead `while True: sleep(60)` and + # NEVER observe the streaming-pull future — so a terminal failure (auth loss, broker + # drop, deleted subscription) would stall this stream silently, invisible to run(). + # Blocking on result() surfaces such a failure as an exception, which propagates to + # run() for a supervised restart (the resilience the other streams get for free). + consumer.stream(block=False) + try: + consumer.streaming_pull_future.result() # blocks; raises on terminal failure + finally: + # Tear the pull/executor down before run() retries, so background threads don't + # leak across restarts. Best-effort — the original failure should propagate. + try: + consumer.stop() + except Exception: + logger.debug(f'{self.STREAM_NAME}: error during consumer.stop()', exc_info=True) + + def normalize_alert(self, raw_alert: pittgoogle.Alert, topic: str = '') -> NormalizedAlert: + """Map a Pitt-Google ZTF alert to a NormalizedAlert. + + Identity comes from the schema-aware accessors (these work for both ztf and lsst): + sourceid -> alert_id the detection that triggered the alert (ZTF candid) + objectid -> object_id the persistent astronomical object (ZTF objectId) + Coordinates, magnitude, and observation time come from the ZTF 'candidate' dict — + alert.ra/dec are None for ZTF, so we read candidate.ra/dec directly. + + Args: + raw_alert: A pittgoogle.Alert (schema_name set by listen()'s callback). + topic: The Pitt-Google topic the alert arrived on. + + Returns: + NormalizedAlert with ZTF fields populated. (LSST diaSource handling is deferred + until Pitt-Google's lsst-* topics go live.) + """ + # Strip cutout stamp data (cutoutScience/Template/Difference) from raw_payload — large, + # binary, compressed FITS that aren't JSON-serializable. Same idiom as Fink. (pittgoogle + # 0.3.22 has no drop_cutouts(); a dict comprehension is predictable and keeps the rest.) + raw_payload = {key: value for key, value in raw_alert.dict.items() if not key.startswith('cutout')} + candidate = raw_payload.get('candidate', {}) # ZTF: nested per-detection measurements + + # published_time = the survey's upstream publish time. Pitt-Google ingests from ZTF's + # Kafka and carries the original Kafka message timestamp (ms epoch) in the Pub/Sub + # attributes as 'kafka.timestamp'. Distinct from the Pub/Sub publishTime that throttles + # ztf-loop (save_ztf_loop_hourly). + kafka_timestamp_ms = raw_alert.attributes.get('kafka.timestamp') + published_time = ( + datetime.fromtimestamp(int(kafka_timestamp_ms) / 1000, tz=timezone.utc) + if kafka_timestamp_ms else None + ) + + # observation_time from the detection's Julian Date. + jd = candidate.get('jd') + observation_time = _jd_to_datetime(jd) if jd is not None else None + + object_id = raw_alert.objectid + return NormalizedAlert( + stream_name=self.STREAM_NAME, + topic=topic, + observation_time=observation_time, + published_time=published_time, + alert_id=str(raw_alert.sourceid), + object_id=str(object_id) if object_id is not None else None, + ra=candidate.get('ra'), + dec=candidate.get('dec'), + magnitude=candidate.get('magpsf'), + flux=None, + raw_payload=raw_payload, + ) diff --git a/tom_alertstreams/apps.py b/tom_alertstreams/apps.py index 824575a..04296dc 100644 --- a/tom_alertstreams/apps.py +++ b/tom_alertstreams/apps.py @@ -1,6 +1,33 @@ from django.apps import AppConfig +from django.conf import settings +from django.urls import include, path class TomAlertstreamsConfig(AppConfig): default_auto_field = 'django.db.models.BigAutoField' name = 'tom_alertstreams' + + def include_url_paths(self) -> list: + """Register tom_alertstreams URLs with the TOM Common URL configuration. + + Returns URL patterns that mount the Recent Alerts page at /alertstreams/, + but only when `SHOW_RECENT_ALERTS = True` is set in settings. If the setting + is absent or False, an empty list is returned and the page is not accessible. + """ + # only return urlpatterns if they opt-in + if not getattr(settings, 'SHOW_RECENT_ALERTS', False): + return [] + return [ + path('alertstreams/', include('tom_alertstreams.urls', namespace='alertstreams')), + ] + + def nav_items(self) -> list: + """Add the 'Recent Alerts' link to the TOM navbar. + + Returns a list of navbar item dicts auto-discovered by the navbar_app_addons + template tag, but only when SHOW_RECENT_ALERTS = True is set in settings. + """ + # only show the navbar item if they opt-in + if not getattr(settings, 'SHOW_RECENT_ALERTS', False): + return [] + return [{'partial': 'tom_alertstreams/partials/navbar_link.html'}] diff --git a/tom_alertstreams/management/commands/readstreams.py b/tom_alertstreams/management/commands/readstreams.py index 4482551..1d6cd65 100644 --- a/tom_alertstreams/management/commands/readstreams.py +++ b/tom_alertstreams/management/commands/readstreams.py @@ -25,14 +25,29 @@ def handle(self, *args, **options): logger.error(f'{ex.__class__.__name__}: Configure alert streams in settings.py ALERT_STREAMS: {ex}') exit(1) + # Run each alert_stream in its own Thread (sort of at the same time). + # Target run() (not listen()) so each stream is supervised: run() catches + # and logs any error from listen() and restarts it, so a single stream's + # broker outage can't silently kill that thread while the others continue. + # daemon=True so a single Ctrl-C (see the join loop below) tears everything down. + threads = [] + for alert_stream in alert_streams: + t = Thread(target=alert_stream.run, name=alert_stream._get_stream_classname(), daemon=True) + t.start() + threads.append(t) + logger.info((f'read_streams {alert_stream._get_stream_classname()} TID={t.native_id} ; ' + f'thread identifier={t.ident}')) + + # Important: Block the main thread on the stream threads instead of returning. + # (i.e. don't let `AlertStream.handle()` return). + + # Why: If AlertStreams.handle() returns, the interpreter enters its + # shutdown/finalize state. In that state, Google Cloud Pub/Sub's gRPC streaming + # pull silently stops delivering messages. So, Pitt-Google, for example, + # streams ingest nothing. Kafka-based streams are not affected. try: - # listen to each alert_stream in it's own Thread (sort of at the same time) - for alert_stream in alert_streams: - t = Thread(target=alert_stream.listen, name=alert_stream._get_stream_classname()) - t.start() - logger.info((f'read_streams {alert_stream._get_stream_classname()} TID={t.native_id} ; ' - f'thread identifier={t.ident}')) - except KeyboardInterrupt as msg: - logger.info(f'read_streams handling KeyboardInterupt {msg}') - - logger.info('readstreams Command.handle() returning...') + while any(thread.is_alive() for thread in threads): # the "join loop" mentioned above + for thread in threads: + thread.join(timeout=1.0) # 1s timeout keeps Ctrl-C working + except KeyboardInterrupt: + logger.info('readstreams: KeyboardInterrupt received; shutting down.') diff --git a/tom_alertstreams/migrations/0001_initial.py b/tom_alertstreams/migrations/0001_initial.py new file mode 100644 index 0000000..b3ddc5c --- /dev/null +++ b/tom_alertstreams/migrations/0001_initial.py @@ -0,0 +1,33 @@ +# Generated by Django 4.2.29 on 2026-03-05 17:25 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + initial = True + + dependencies = [ + ] + + operations = [ + migrations.CreateModel( + name='Alert', + fields=[ + ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('stream_name', models.CharField(db_index=True, max_length=100)), + ('topic', models.CharField(max_length=200)), + ('timestamp', models.DateTimeField(db_index=True)), + ('alert_id', models.CharField(max_length=200)), + ('object_id', models.CharField(blank=True, max_length=200, null=True)), + ('ra', models.FloatField(null=True)), + ('dec', models.FloatField(null=True)), + ('magnitude', models.FloatField(null=True)), + ('raw_payload', models.JSONField(default=dict)), + ], + options={ + 'ordering': ['-timestamp'], + 'indexes': [models.Index(fields=['stream_name', 'timestamp'], name='tom_alertst_stream__a592d5_idx')], + }, + ), + ] diff --git a/tom_alertstreams/migrations/0002_alert_flux.py b/tom_alertstreams/migrations/0002_alert_flux.py new file mode 100644 index 0000000..34d9f99 --- /dev/null +++ b/tom_alertstreams/migrations/0002_alert_flux.py @@ -0,0 +1,18 @@ +# Generated by Django 4.2.29 on 2026-03-11 00:21 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('tom_alertstreams', '0001_initial'), + ] + + operations = [ + migrations.AddField( + model_name='alert', + name='flux', + field=models.FloatField(null=True), + ), + ] diff --git a/tom_alertstreams/migrations/0003_remove_alert_tom_alertst_stream__a592d5_idx_and_more.py b/tom_alertstreams/migrations/0003_remove_alert_tom_alertst_stream__a592d5_idx_and_more.py new file mode 100644 index 0000000..9956625 --- /dev/null +++ b/tom_alertstreams/migrations/0003_remove_alert_tom_alertst_stream__a592d5_idx_and_more.py @@ -0,0 +1,21 @@ +# Generated by Django 4.2.29 on 2026-03-16 23:51 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('tom_alertstreams', '0002_alert_flux'), + ] + + operations = [ + migrations.RemoveIndex( + model_name='alert', + name='tom_alertst_stream__a592d5_idx', + ), + migrations.AddIndex( + model_name='alert', + index=models.Index(fields=['stream_name', 'topic', 'timestamp'], name='tom_alertst_stream__a32a9e_idx'), + ), + ] diff --git a/tom_alertstreams/migrations/0004_alert_created.py b/tom_alertstreams/migrations/0004_alert_created.py new file mode 100644 index 0000000..b86219e --- /dev/null +++ b/tom_alertstreams/migrations/0004_alert_created.py @@ -0,0 +1,18 @@ +# Generated by Django 4.2.29 on 2026-06-05 01:45 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('tom_alertstreams', '0003_remove_alert_tom_alertst_stream__a592d5_idx_and_more'), + ] + + operations = [ + migrations.AddField( + model_name='alert', + name='created', + field=models.DateTimeField(auto_now_add=True, db_index=True, null=True), + ), + ] diff --git a/tom_alertstreams/migrations/0005_alter_alert_options.py b/tom_alertstreams/migrations/0005_alter_alert_options.py new file mode 100644 index 0000000..47a25f8 --- /dev/null +++ b/tom_alertstreams/migrations/0005_alter_alert_options.py @@ -0,0 +1,17 @@ +# Generated by Django 4.2.29 on 2026-06-09 17:38 + +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('tom_alertstreams', '0004_alert_created'), + ] + + operations = [ + migrations.AlterModelOptions( + name='alert', + options={'ordering': ['-created']}, + ), + ] diff --git a/tom_alertstreams/migrations/0006_remove_alert_tom_alertst_stream__a32a9e_idx_and_more.py b/tom_alertstreams/migrations/0006_remove_alert_tom_alertst_stream__a32a9e_idx_and_more.py new file mode 100644 index 0000000..c326382 --- /dev/null +++ b/tom_alertstreams/migrations/0006_remove_alert_tom_alertst_stream__a32a9e_idx_and_more.py @@ -0,0 +1,35 @@ +# Generated by Django 4.2.29 on 2026-06-09 22:19 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('tom_alertstreams', '0005_alter_alert_options'), + ] + + operations = [ + migrations.RemoveIndex( + model_name='alert', + name='tom_alertst_stream__a32a9e_idx', + ), + migrations.RemoveField( + model_name='alert', + name='timestamp', + ), + migrations.AddField( + model_name='alert', + name='observation_time', + field=models.DateTimeField(db_index=True, null=True), + ), + migrations.AddField( + model_name='alert', + name='published_time', + field=models.DateTimeField(null=True), + ), + migrations.AddIndex( + model_name='alert', + index=models.Index(fields=['stream_name', 'topic', 'observation_time'], name='tom_alertst_stream__a16eaf_idx'), + ), + ] diff --git a/tom_alertstreams/models.py b/tom_alertstreams/models.py index 71a8362..b7ff301 100644 --- a/tom_alertstreams/models.py +++ b/tom_alertstreams/models.py @@ -1,3 +1,130 @@ +from __future__ import annotations + +from typing import Any, ClassVar + +from django.conf import settings from django.db import models -# Create your models here. + +class FIFOQueueMixin(models.Model): + """Mixin that enforces a per-partition maximum row count. + + Turns a model's table into a First-In-First-Out (FIFO) queue by extending + save() to delete the oldest rows beyond FIFO_MAX after every insert. + + Subclasses set FIFO_MAX and FIFO_PARTITION_FIELDS as class variables. + + Partitioning: FIFO_PARTITION_FIELDS names one or more model fields whose + values together define a partition. Each unique combination of partition + field values gets its own independent FIFO budget of FIFO_MAX rows. + + With a single partition field (e.g. ``('stream_name',)``), the table holds + at most ``FIFO_MAX × number_of_distinct_stream_names`` rows. With compound + fields (e.g. ``('stream_name', 'topic')``), each unique combination gets + its own budget, so the table can hold up to + ``FIFO_MAX × number_of_distinct_(stream_name, topic)_pairs`` rows. + For example: 8 streams × 19 topics/stream × 1000 max = 152,000 rows + (though in practice most streams have far fewer topics). + + If FIFO_PARTITION_FIELDS is None, the limit applies to the entire table. + """ + FIFO_MAX: ClassVar[int] = 10 + FIFO_PARTITION_FIELDS: ClassVar[tuple[str, ...] | None] = None + + class Meta: + abstract = True + + def save(self, *args: Any, **kwargs: Any) -> None: + super().save(*args, **kwargs) # save, then trim: Ensures the newly-saved + # this is the extention to the method: + self._enforce_fifo_limit() # row is counted against the FIFO_MAX limit. + + def _enforce_fifo_limit(self) -> None: + """Delete rows beyond FIFO_MAX, oldest-received first, within this partition. + + The partition is defined by FIFO_PARTITION_FIELDS — all rows sharing the + same values across those fields form one partition. Each partition is + independently capped at FIFO_MAX rows. + + Eviction is by `created` (insertion/received time) — true FIFO order — NOT by + the alert's observation `timestamp`. Ordering a FIFO by observation time would + instantly evict a freshly-received alert whose observation time is old (e.g. a + broker streaming last night's alerts now, after the survey went quiet), so it + would never appear in the table. By `created`, the most-recently-received rows + are always kept. (The concrete model must provide a `created` insertion-time + field; Alert does, via auto_now_add.) + + Uses list() to materialise PKs before the DELETE to avoid a SQLite restriction + that forbids DELETE from a table referenced in the same statement's subquery. + """ + qs = self.__class__.objects.all() # query set + if self.FIFO_PARTITION_FIELDS is not None: + # Scope to rows sharing this instance's values for all partition fields + partition_filter = { + field: getattr(self, field) for field in self.FIFO_PARTITION_FIELDS + } + qs = qs.filter(**partition_filter) + # Keep the FIFO_MAX most-recently-received rows; materialise PKs first to avoid + # a SQLite subquery-in-DELETE restriction. + excess_pks = list( + # the list of the pks beyond FIFO_MAX + qs.order_by('-created').values_list('pk', flat=True)[self.FIFO_MAX:] + ) + if excess_pks: + self.__class__.objects.filter(pk__in=excess_pks).delete() # delete overflow + + +class Alert(FIFOQueueMixin): + """A normalized alert received from an alert stream, stored for recent display. + + This class is designed specifically for a Recent Alerts demonstration page. + It's the model shown in the table of recent alerts and size is limited by + the FIFOQueueMixin, FIFO_MAX, and FIFO_PARTITION_FIELDS. + """ + # Set the mixin class variables — partition by (stream_name, topic) so each + # stream+topic combination gets its own independent FIFO budget. + FIFO_MAX: ClassVar[int] = getattr(settings, 'ALERTSTREAMS_RECENT_COUNT', 10) + FIFO_PARTITION_FIELDS: ClassVar[tuple[str, ...] | None] = ('stream_name', 'topic') + + stream_name = models.CharField(max_length=100, db_index=True) + topic = models.CharField(max_length=200) + + # When the telescope observed the source the alert is about (a detection time). + # Nullable: some alerts (e.g. GCN Circulars) are not about a single observation. + observation_time = models.DateTimeField(null=True, db_index=True) + + # When the broker/survey issued/published the alert (e.g. GCN Circular createdOn, + # Fink brokerEndProcessTimestamp). Nullable: not every stream exposes it. The gap + # between observation_time and published_time is the broker's processing latency. + published_time = models.DateTimeField(null=True) + + alert_id = models.CharField(max_length=200) + object_id = models.CharField(max_length=200, blank=True, null=True) + ra = models.FloatField(null=True) + dec = models.FloatField(null=True) + magnitude = models.FloatField(null=True) + flux = models.FloatField(null=True) + + raw_payload = models.JSONField(default=dict) + + # When the row was received and saved — our receipt clock, distinct from + # `observation_time` (when observed) and `published_time` (when the broker issued + # it). Observation/publish times can be stale or backlogged, so `created` is what + # answers "is this stream still ingesting?" — and it drives ordering + the FIFO. + # Nullable so the column could be added while a running readstreams (old model) was + # still inserting; every new row sets it via auto_now_add. + created = models.DateTimeField(auto_now_add=True, null=True, db_index=True) + + class Meta(FIFOQueueMixin.Meta): # this is the way you subclass the internal Meta class + + # in the FIFOQueueMixin.Meta, abstract is True, + abstract = False # here, in the concrete model, we're not abstract. + + ordering = ['-created'] # most recently received + + # speed up the HTMX filtering + indexes = [models.Index(fields=['stream_name', 'topic', 'observation_time'])] + + def __str__(self) -> str: + return (f'Alert {self.alert_id} received from {self.stream_name} on topic {self.topic} ' + f'at {self.created}') diff --git a/tom_alertstreams/tables.py b/tom_alertstreams/tables.py new file mode 100644 index 0000000..f513199 --- /dev/null +++ b/tom_alertstreams/tables.py @@ -0,0 +1,385 @@ +from __future__ import annotations + +import logging +import urllib.parse +from typing import Any, ClassVar + +from crispy_forms.layout import Column, Layout, Row +import django_filters +import django_tables2 as tables +from django import forms +from django.db.models import Q +from django.utils.html import format_html + +from tom_alertstreams.alertstreams.alertstream import get_alert_stream_classes +from tom_alertstreams.models import Alert +from tom_common.htmx_table import HTMXTable, HTMXTableFilterSet + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# AlertTable +# --------------------------------------------------------------------------- + +class AlertTable(HTMXTable): + """HTMX-driven table of recent alerts from all configured alert streams. + + Receives a presenter_map at construction time — a dict mapping stream_name + to an AlertStreamPresenter instance. render_alert_id() and render_object_id() + delegate URL construction to the presenter, keeping this table fully generic + with zero stream-specific logic. + """ + # Custom partial that includes an OOB swap to update the stream status dashboard + partial_template_name = 'tom_alertstreams/partials/alert_table_partial.html' + + def __init__( + self, + *args: Any, + presenter_map: dict[str, AlertStreamPresenter] | None = None, + **kwargs: Any, + ) -> None: + self.presenter_map = presenter_map or {} + super().__init__(*args, **kwargs) + + # + # Custom field renderers + # + + stream_name = tables.Column(verbose_name='Stream') # sets the column header value + # Two-word headers so these time columns wrap rather than stretch the table wide. + observation_time = tables.Column(verbose_name='Observation Time') + published_time = tables.Column(verbose_name='Published Time') + + # render_FIELDNAME() methods are called automatically when present + def render_observation_time(self, value: Any) -> str: + """Render the observation time in UTC 24-hour format (blank if absent). + + The result looks like this: 2026-03-05 18:51:30 UTC. Blank for alerts with no + observation (e.g. GCN Circulars). + """ + return value.strftime('%Y-%m-%d %H:%M:%S UTC') if value else '' + + def render_published_time(self, value: Any) -> str: + """Render the broker/survey publish time in UTC 24-hour format (blank if absent).""" + return value.strftime('%Y-%m-%d %H:%M:%S UTC') if value else '' + + def render_alert_id(self, record: Alert, value: str) -> str: + """Render alert_id as a hyperlink if the stream's presenter provides a URL.""" + presenter = self.presenter_map.get(record.stream_name) + if presenter: + url = presenter.alert_url(record) + if url: + return format_html('{}', url, value) + return value + + def render_object_id(self, record: Alert, value: str) -> str: + """Render object_id as a hyperlink if the stream's presenter provides a URL.""" + if not value: + return value + presenter = self.presenter_map.get(record.stream_name) + if presenter: + url = presenter.object_url(record) + if url: + return format_html('{}', url, value) + return value + + def render_ra(self, value: Any) -> str: + """Render RA to 5 decimal places (~0.04 arcsec, matching LSST precision).""" + return f'{value:.5f}' if value is not None else '' + + def render_dec(self, value: Any) -> str: + """Render Dec to 5 decimal places (~0.04 arcsec, matching LSST precision).""" + return f'{value:.5f}' if value is not None else '' + + def render_magnitude(self, value: Any) -> str: + """Render magnitude to 3 decimal places (~1 mmag, matching survey photometric precision).""" + return f'{value:.3f}' if value is not None else '' + + def render_flux(self, value: Any) -> str: + """Render flux in nanojansky to 3 decimal places.""" + return f'{value:.3f}' if value is not None else '' + + class Meta(HTMXTable.Meta): + model = Alert + fields = [ + 'selection', 'alert_id', 'stream_name', 'topic', 'published_time', + 'observation_time', 'object_id', 'ra', 'dec', 'magnitude', 'flux', + ] + + +def _get_stream_name_choices() -> list[tuple[str, str]]: + """Build dropdown choices from the active streams in settings.ALERT_STREAMS. + + Called at form-render time (not import time) so changes to settings take + effect without restarting the process. Returns a list of (value, label) + tuples using each stream's STREAM_NAME. Streams that fail to import are + silently skipped so a misconfigured entry doesn't break the filter form. + + Deduped by STREAM_NAME: a broker may be configured as several ALERT_STREAMS + entries that share a name (Pitt-Google runs one entry per Pub/Sub topic, all + 'pittgoogle'), and the dropdown should show one choice per stream. + """ + seen: set[str] = set() + choices: list[tuple[str, str]] = [] + for klass in get_alert_stream_classes(): + if klass.STREAM_NAME in seen: + continue + seen.add(klass.STREAM_NAME) + choices.append((klass.STREAM_NAME, klass.STREAM_NAME)) + return choices + + +class AlertFilterSet(HTMXTableFilterSet): + """FilterSet for the Recent Alerts table. + + Provides a 'query' full-text search (inherited from HTMXTableFilterSet) plus + the fields defined here, which appear in the Advanced> expansion of the form. + + Implements the HTMX "cascading select" pattern: the topic dropdown choices + are scoped to the currently selected stream. See __init__ for the dynamic + choice logic and recent_alerts.html for the HTMX trigger that refreshes + the topic + innerHTML. When no stream is selected, returns topics from all streams. + """ + stream_name = request.GET.get('stream_name', '') + qs = Alert.objects.all() + if stream_name: + qs = qs.filter(stream_name=stream_name) + topics = qs.values_list('topic', flat=True).distinct().order_by('topic') + + # Build '] + options.extend(f'' for topic in topics) + return HttpResponse('\n'.join(options)) diff --git a/typings/babamul/__init__.pyi b/typings/babamul/__init__.pyi new file mode 100644 index 0000000..6b6ea3b --- /dev/null +++ b/typings/babamul/__init__.pyi @@ -0,0 +1,9 @@ +"""Type stubs for babamul. + +Babamul does not ship a py.typed marker, so Pylance/Pyright cannot resolve +its types. These minimal stubs cover only what tom_alertstreams imports. +Remove this directory when babamul adds py.typed to its distribution. +""" + +from .consumer import AlertConsumer as AlertConsumer +from .models import LsstAlert as LsstAlert, ZtfAlert as ZtfAlert diff --git a/typings/babamul/consumer.pyi b/typings/babamul/consumer.pyi new file mode 100644 index 0000000..500ce9a --- /dev/null +++ b/typings/babamul/consumer.pyi @@ -0,0 +1,28 @@ +"""Type stub for babamul.consumer.""" + +from collections.abc import Iterator +from typing import Any + +from .models import LsstAlert, ZtfAlert + +class AlertConsumer: + def __init__( + self, + topics: str | list[str] = ..., + username: str | None = ..., + password: str | None = ..., + server: str | None = ..., + group_id: str | None = ..., + offset: str = ..., + timeout: float | None = ..., + auto_commit: bool = ..., + as_raw: bool = ..., + ) -> None: ... + def __iter__(self) -> Iterator[ZtfAlert | LsstAlert | dict[str, Any]]: ... + def __enter__(self) -> AlertConsumer: ... + def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object) -> None: ... + def close(self) -> None: ... + @property + def topics(self) -> list[str]: ... + @property + def group_id(self) -> str: ... diff --git a/typings/babamul/models.pyi b/typings/babamul/models.pyi new file mode 100644 index 0000000..8af70e3 --- /dev/null +++ b/typings/babamul/models.pyi @@ -0,0 +1,36 @@ +"""Type stub for babamul.models — covers only what tom_alertstreams uses.""" + +from datetime import datetime + +from pydantic import BaseModel + +class ZtfCandidate(BaseModel): + jd: float + ra: float + dec: float + magpsf: float + sigmapsf: float + @property + def datetime(self) -> datetime: ... + +class LsstCandidate(BaseModel): + ra: float + dec: float + magpsf: float + sigmapsf: float + jd: float + objectId: str + @property + def datetime(self) -> datetime: ... + +class ZtfAlert(BaseModel): + candid: int + objectId: str + candidate: ZtfCandidate + topic: str | None + +class LsstAlert(BaseModel): + candid: int + objectId: str + candidate: LsstCandidate + topic: str | None