Vendor backend and frontend into the monorepo
The submodule collapse was only half applied: .gitmodules was deleted but backend/ and frontend/ were still recorded as gitlinks, so none of their files were tracked. Replace the gitlinks with the real trees. Also untrack .env (it carried placeholder secrets) in favour of a tracked .env.example, drop the committed __pycache__, and narrow the blanket *.png ignore that would have swallowed design assets. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
3cb16958ba
commit
1916f7f778
Executable
+1
@@ -0,0 +1 @@
|
||||
Generic single-database configuration.
|
||||
Executable
+84
@@ -0,0 +1,84 @@
|
||||
import os
|
||||
from logging.config import fileConfig
|
||||
|
||||
from alembic import context
|
||||
from sqlalchemy import engine_from_config, pool
|
||||
|
||||
# this is the Alembic Config object, which provides
|
||||
# access to the values within the .ini file in use.
|
||||
config = context.config
|
||||
|
||||
# Interpret the config file for Python logging.
|
||||
# This line sets up loggers basically.
|
||||
fileConfig(config.config_file_name)
|
||||
|
||||
# add your model's MetaData object here
|
||||
# for 'autogenerate' support
|
||||
# from myapp import mymodel
|
||||
# target_metadata = mymodel.Base.metadata
|
||||
# target_metadata = None
|
||||
|
||||
from app.models import SQLModel # noqa
|
||||
from app.core.config import settings # noqa
|
||||
|
||||
target_metadata = SQLModel.metadata
|
||||
|
||||
# other values from the config, defined by the needs of env.py,
|
||||
# can be acquired:
|
||||
# my_important_option = config.get_main_option("my_important_option")
|
||||
# ... etc.
|
||||
|
||||
|
||||
def get_url():
|
||||
return str(settings.SQLALCHEMY_DATABASE_URI)
|
||||
|
||||
|
||||
def run_migrations_offline():
|
||||
"""Run migrations in 'offline' mode.
|
||||
|
||||
This configures the context with just a URL
|
||||
and not an Engine, though an Engine is acceptable
|
||||
here as well. By skipping the Engine creation
|
||||
we don't even need a DBAPI to be available.
|
||||
|
||||
Calls to context.execute() here emit the given string to the
|
||||
script output.
|
||||
|
||||
"""
|
||||
url = get_url()
|
||||
context.configure(
|
||||
url=url, target_metadata=target_metadata, literal_binds=True, compare_type=True
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
def run_migrations_online():
|
||||
"""Run migrations in 'online' mode.
|
||||
|
||||
In this scenario we need to create an Engine
|
||||
and associate a connection with the context.
|
||||
|
||||
"""
|
||||
configuration = config.get_section(config.config_ini_section)
|
||||
configuration["sqlalchemy.url"] = get_url()
|
||||
connectable = engine_from_config(
|
||||
configuration,
|
||||
prefix="sqlalchemy.",
|
||||
poolclass=pool.NullPool,
|
||||
)
|
||||
|
||||
with connectable.connect() as connection:
|
||||
context.configure(
|
||||
connection=connection, target_metadata=target_metadata, compare_type=True
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
if context.is_offline_mode():
|
||||
run_migrations_offline()
|
||||
else:
|
||||
run_migrations_online()
|
||||
Executable
+25
@@ -0,0 +1,25 @@
|
||||
"""${message}
|
||||
|
||||
Revision ID: ${up_revision}
|
||||
Revises: ${down_revision | comma,n}
|
||||
Create Date: ${create_date}
|
||||
|
||||
"""
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel.sql.sqltypes
|
||||
${imports if imports else ""}
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = ${repr(up_revision)}
|
||||
down_revision = ${repr(down_revision)}
|
||||
branch_labels = ${repr(branch_labels)}
|
||||
depends_on = ${repr(depends_on)}
|
||||
|
||||
|
||||
def upgrade():
|
||||
${upgrades if upgrades else "pass"}
|
||||
|
||||
|
||||
def downgrade():
|
||||
${downgrades if downgrades else "pass"}
|
||||
Executable
@@ -0,0 +1,37 @@
|
||||
"""Add cascade delete relationships
|
||||
|
||||
Revision ID: 1a31ce608336
|
||||
Revises: d98dd8ec85a3
|
||||
Create Date: 2024-07-31 22:24:34.447891
|
||||
|
||||
"""
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel.sql.sqltypes
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = '1a31ce608336'
|
||||
down_revision = 'd98dd8ec85a3'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.alter_column('item', 'owner_id',
|
||||
existing_type=sa.UUID(),
|
||||
nullable=False)
|
||||
op.drop_constraint('item_owner_id_fkey', 'item', type_='foreignkey')
|
||||
op.create_foreign_key(None, 'item', 'user', ['owner_id'], ['id'], ondelete='CASCADE')
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade():
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_constraint(None, 'item', type_='foreignkey')
|
||||
op.create_foreign_key('item_owner_id_fkey', 'item', 'user', ['owner_id'], ['id'])
|
||||
op.alter_column('item', 'owner_id',
|
||||
existing_type=sa.UUID(),
|
||||
nullable=True)
|
||||
# ### end Alembic commands ###
|
||||
+69
@@ -0,0 +1,69 @@
|
||||
"""Add max length for string(varchar) fields in User and Items models
|
||||
|
||||
Revision ID: 9c0a54914c78
|
||||
Revises: e2412789c190
|
||||
Create Date: 2024-06-17 14:42:44.639457
|
||||
|
||||
"""
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel.sql.sqltypes
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = '9c0a54914c78'
|
||||
down_revision = 'e2412789c190'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
# Adjust the length of the email field in the User table
|
||||
op.alter_column('user', 'email',
|
||||
existing_type=sa.String(),
|
||||
type_=sa.String(length=255),
|
||||
existing_nullable=False)
|
||||
|
||||
# Adjust the length of the full_name field in the User table
|
||||
op.alter_column('user', 'full_name',
|
||||
existing_type=sa.String(),
|
||||
type_=sa.String(length=255),
|
||||
existing_nullable=True)
|
||||
|
||||
# Adjust the length of the title field in the Item table
|
||||
op.alter_column('item', 'title',
|
||||
existing_type=sa.String(),
|
||||
type_=sa.String(length=255),
|
||||
existing_nullable=False)
|
||||
|
||||
# Adjust the length of the description field in the Item table
|
||||
op.alter_column('item', 'description',
|
||||
existing_type=sa.String(),
|
||||
type_=sa.String(length=255),
|
||||
existing_nullable=True)
|
||||
|
||||
|
||||
def downgrade():
|
||||
# Revert the length of the email field in the User table
|
||||
op.alter_column('user', 'email',
|
||||
existing_type=sa.String(length=255),
|
||||
type_=sa.String(),
|
||||
existing_nullable=False)
|
||||
|
||||
# Revert the length of the full_name field in the User table
|
||||
op.alter_column('user', 'full_name',
|
||||
existing_type=sa.String(length=255),
|
||||
type_=sa.String(),
|
||||
existing_nullable=True)
|
||||
|
||||
# Revert the length of the title field in the Item table
|
||||
op.alter_column('item', 'title',
|
||||
existing_type=sa.String(length=255),
|
||||
type_=sa.String(),
|
||||
existing_nullable=False)
|
||||
|
||||
# Revert the length of the description field in the Item table
|
||||
op.alter_column('item', 'description',
|
||||
existing_type=sa.String(length=255),
|
||||
type_=sa.String(),
|
||||
existing_nullable=True)
|
||||
+90
@@ -0,0 +1,90 @@
|
||||
"""Edit replace id integers in all models to use UUID instead
|
||||
|
||||
Revision ID: d98dd8ec85a3
|
||||
Revises: 9c0a54914c78
|
||||
Create Date: 2024-07-19 04:08:04.000976
|
||||
|
||||
"""
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel.sql.sqltypes
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = 'd98dd8ec85a3'
|
||||
down_revision = '9c0a54914c78'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
# Ensure uuid-ossp extension is available
|
||||
op.execute('CREATE EXTENSION IF NOT EXISTS "uuid-ossp"')
|
||||
|
||||
# Create a new UUID column with a default UUID value
|
||||
op.add_column('user', sa.Column('new_id', postgresql.UUID(as_uuid=True), default=sa.text('uuid_generate_v4()')))
|
||||
op.add_column('item', sa.Column('new_id', postgresql.UUID(as_uuid=True), default=sa.text('uuid_generate_v4()')))
|
||||
op.add_column('item', sa.Column('new_owner_id', postgresql.UUID(as_uuid=True), nullable=True))
|
||||
|
||||
# Populate the new columns with UUIDs
|
||||
op.execute('UPDATE "user" SET new_id = uuid_generate_v4()')
|
||||
op.execute('UPDATE item SET new_id = uuid_generate_v4()')
|
||||
op.execute('UPDATE item SET new_owner_id = (SELECT new_id FROM "user" WHERE "user".id = item.owner_id)')
|
||||
|
||||
# Set the new_id as not nullable
|
||||
op.alter_column('user', 'new_id', nullable=False)
|
||||
op.alter_column('item', 'new_id', nullable=False)
|
||||
|
||||
# Drop old columns and rename new columns
|
||||
op.drop_constraint('item_owner_id_fkey', 'item', type_='foreignkey')
|
||||
op.drop_column('item', 'owner_id')
|
||||
op.alter_column('item', 'new_owner_id', new_column_name='owner_id')
|
||||
|
||||
op.drop_column('user', 'id')
|
||||
op.alter_column('user', 'new_id', new_column_name='id')
|
||||
|
||||
op.drop_column('item', 'id')
|
||||
op.alter_column('item', 'new_id', new_column_name='id')
|
||||
|
||||
# Create primary key constraint
|
||||
op.create_primary_key('user_pkey', 'user', ['id'])
|
||||
op.create_primary_key('item_pkey', 'item', ['id'])
|
||||
|
||||
# Recreate foreign key constraint
|
||||
op.create_foreign_key('item_owner_id_fkey', 'item', 'user', ['owner_id'], ['id'])
|
||||
|
||||
def downgrade():
|
||||
# Reverse the upgrade process
|
||||
op.add_column('user', sa.Column('old_id', sa.Integer, autoincrement=True))
|
||||
op.add_column('item', sa.Column('old_id', sa.Integer, autoincrement=True))
|
||||
op.add_column('item', sa.Column('old_owner_id', sa.Integer, nullable=True))
|
||||
|
||||
# Populate the old columns with default values
|
||||
# Generate sequences for the integer IDs if not exist
|
||||
op.execute('CREATE SEQUENCE IF NOT EXISTS user_id_seq AS INTEGER OWNED BY "user".old_id')
|
||||
op.execute('CREATE SEQUENCE IF NOT EXISTS item_id_seq AS INTEGER OWNED BY item.old_id')
|
||||
|
||||
op.execute('SELECT setval(\'user_id_seq\', COALESCE((SELECT MAX(old_id) + 1 FROM "user"), 1), false)')
|
||||
op.execute('SELECT setval(\'item_id_seq\', COALESCE((SELECT MAX(old_id) + 1 FROM item), 1), false)')
|
||||
|
||||
op.execute('UPDATE "user" SET old_id = nextval(\'user_id_seq\')')
|
||||
op.execute('UPDATE item SET old_id = nextval(\'item_id_seq\'), old_owner_id = (SELECT old_id FROM "user" WHERE "user".id = item.owner_id)')
|
||||
|
||||
# Drop new columns and rename old columns back
|
||||
op.drop_constraint('item_owner_id_fkey', 'item', type_='foreignkey')
|
||||
op.drop_column('item', 'owner_id')
|
||||
op.alter_column('item', 'old_owner_id', new_column_name='owner_id')
|
||||
|
||||
op.drop_column('user', 'id')
|
||||
op.alter_column('user', 'old_id', new_column_name='id')
|
||||
|
||||
op.drop_column('item', 'id')
|
||||
op.alter_column('item', 'old_id', new_column_name='id')
|
||||
|
||||
# Create primary key constraint
|
||||
op.create_primary_key('user_pkey', 'user', ['id'])
|
||||
op.create_primary_key('item_pkey', 'item', ['id'])
|
||||
|
||||
# Recreate foreign key constraint
|
||||
op.create_foreign_key('item_owner_id_fkey', 'item', 'user', ['owner_id'], ['id'])
|
||||
@@ -0,0 +1,54 @@
|
||||
"""Initialize models
|
||||
|
||||
Revision ID: e2412789c190
|
||||
Revises:
|
||||
Create Date: 2023-11-24 22:55:43.195942
|
||||
|
||||
"""
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel.sql.sqltypes
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "e2412789c190"
|
||||
down_revision = None
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.create_table(
|
||||
"user",
|
||||
sa.Column("email", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
|
||||
sa.Column("is_active", sa.Boolean(), nullable=False),
|
||||
sa.Column("is_superuser", sa.Boolean(), nullable=False),
|
||||
sa.Column("full_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column(
|
||||
"hashed_password", sqlmodel.sql.sqltypes.AutoString(), nullable=False
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
op.create_index(op.f("ix_user_email"), "user", ["email"], unique=True)
|
||||
op.create_table(
|
||||
"item",
|
||||
sa.Column("description", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("title", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
|
||||
sa.Column("owner_id", sa.Integer(), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["owner_id"],
|
||||
["user.id"],
|
||||
),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade():
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_table("item")
|
||||
op.drop_index(op.f("ix_user_email"), table_name="user")
|
||||
op.drop_table("user")
|
||||
# ### end Alembic commands ###
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Add created_at to User and Item
|
||||
|
||||
Revision ID: fe56fa70289e
|
||||
Revises: 1a31ce608336
|
||||
Create Date: 2026-01-23 15:50:37.171462
|
||||
|
||||
"""
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel.sql.sqltypes
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = 'fe56fa70289e'
|
||||
down_revision = '1a31ce608336'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.add_column('item', sa.Column('created_at', sa.DateTime(timezone=True), nullable=True))
|
||||
op.add_column('user', sa.Column('created_at', sa.DateTime(timezone=True), nullable=True))
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
||||
def downgrade():
|
||||
# ### commands auto generated by Alembic - please adjust! ###
|
||||
op.drop_column('user', 'created_at')
|
||||
op.drop_column('item', 'created_at')
|
||||
# ### end Alembic commands ###
|
||||
@@ -0,0 +1,57 @@
|
||||
from collections.abc import Generator
|
||||
from typing import Annotated
|
||||
|
||||
import jwt
|
||||
from fastapi import Depends, HTTPException, status
|
||||
from fastapi.security import OAuth2PasswordBearer
|
||||
from jwt.exceptions import InvalidTokenError
|
||||
from pydantic import ValidationError
|
||||
from sqlmodel import Session
|
||||
|
||||
from app.core import security
|
||||
from app.core.config import settings
|
||||
from app.core.db import engine
|
||||
from app.models import TokenPayload, User
|
||||
|
||||
reusable_oauth2 = OAuth2PasswordBearer(
|
||||
tokenUrl=f"{settings.API_V1_STR}/login/access-token"
|
||||
)
|
||||
|
||||
|
||||
def get_db() -> Generator[Session, None, None]:
|
||||
with Session(engine) as session:
|
||||
yield session
|
||||
|
||||
|
||||
SessionDep = Annotated[Session, Depends(get_db)]
|
||||
TokenDep = Annotated[str, Depends(reusable_oauth2)]
|
||||
|
||||
|
||||
def get_current_user(session: SessionDep, token: TokenDep) -> User:
|
||||
try:
|
||||
payload = jwt.decode(
|
||||
token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
|
||||
)
|
||||
token_data = TokenPayload(**payload)
|
||||
except (InvalidTokenError, ValidationError):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Could not validate credentials",
|
||||
)
|
||||
user = session.get(User, token_data.sub)
|
||||
if not user:
|
||||
raise HTTPException(status_code=404, detail="User not found")
|
||||
if not user.is_active:
|
||||
raise HTTPException(status_code=400, detail="Inactive user")
|
||||
return user
|
||||
|
||||
|
||||
CurrentUser = Annotated[User, Depends(get_current_user)]
|
||||
|
||||
|
||||
def get_current_active_superuser(current_user: CurrentUser) -> User:
|
||||
if not current_user.is_superuser:
|
||||
raise HTTPException(
|
||||
status_code=403, detail="The user doesn't have enough privileges"
|
||||
)
|
||||
return current_user
|
||||
@@ -0,0 +1,14 @@
|
||||
from fastapi import APIRouter
|
||||
|
||||
from app.api.routes import items, login, private, users, utils
|
||||
from app.core.config import settings
|
||||
|
||||
api_router = APIRouter()
|
||||
api_router.include_router(login.router)
|
||||
api_router.include_router(users.router)
|
||||
api_router.include_router(utils.router)
|
||||
api_router.include_router(items.router)
|
||||
|
||||
|
||||
if settings.ENVIRONMENT == "local":
|
||||
api_router.include_router(private.router)
|
||||
@@ -0,0 +1,112 @@
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from sqlmodel import func, select
|
||||
|
||||
from app.api.deps import CurrentUser, SessionDep
|
||||
from app.models import Item, ItemCreate, ItemPublic, ItemsPublic, ItemUpdate, Message
|
||||
|
||||
router = APIRouter(prefix="/items", tags=["items"])
|
||||
|
||||
|
||||
@router.get("/", response_model=ItemsPublic)
|
||||
def read_items(
|
||||
session: SessionDep, current_user: CurrentUser, skip: int = 0, limit: int = 100
|
||||
) -> Any:
|
||||
"""
|
||||
Retrieve items.
|
||||
"""
|
||||
|
||||
if current_user.is_superuser:
|
||||
count_statement = select(func.count()).select_from(Item)
|
||||
count = session.exec(count_statement).one()
|
||||
statement = (
|
||||
select(Item).order_by(Item.created_at.desc()).offset(skip).limit(limit)
|
||||
)
|
||||
items = session.exec(statement).all()
|
||||
else:
|
||||
count_statement = (
|
||||
select(func.count())
|
||||
.select_from(Item)
|
||||
.where(Item.owner_id == current_user.id)
|
||||
)
|
||||
count = session.exec(count_statement).one()
|
||||
statement = (
|
||||
select(Item)
|
||||
.where(Item.owner_id == current_user.id)
|
||||
.order_by(Item.created_at.desc())
|
||||
.offset(skip)
|
||||
.limit(limit)
|
||||
)
|
||||
items = session.exec(statement).all()
|
||||
|
||||
return ItemsPublic(data=items, count=count)
|
||||
|
||||
|
||||
@router.get("/{id}", response_model=ItemPublic)
|
||||
def read_item(session: SessionDep, current_user: CurrentUser, id: uuid.UUID) -> Any:
|
||||
"""
|
||||
Get item by ID.
|
||||
"""
|
||||
item = session.get(Item, id)
|
||||
if not item:
|
||||
raise HTTPException(status_code=404, detail="Item not found")
|
||||
if not current_user.is_superuser and (item.owner_id != current_user.id):
|
||||
raise HTTPException(status_code=403, detail="Not enough permissions")
|
||||
return item
|
||||
|
||||
|
||||
@router.post("/", response_model=ItemPublic)
|
||||
def create_item(
|
||||
*, session: SessionDep, current_user: CurrentUser, item_in: ItemCreate
|
||||
) -> Any:
|
||||
"""
|
||||
Create new item.
|
||||
"""
|
||||
item = Item.model_validate(item_in, update={"owner_id": current_user.id})
|
||||
session.add(item)
|
||||
session.commit()
|
||||
session.refresh(item)
|
||||
return item
|
||||
|
||||
|
||||
@router.put("/{id}", response_model=ItemPublic)
|
||||
def update_item(
|
||||
*,
|
||||
session: SessionDep,
|
||||
current_user: CurrentUser,
|
||||
id: uuid.UUID,
|
||||
item_in: ItemUpdate,
|
||||
) -> Any:
|
||||
"""
|
||||
Update an item.
|
||||
"""
|
||||
item = session.get(Item, id)
|
||||
if not item:
|
||||
raise HTTPException(status_code=404, detail="Item not found")
|
||||
if not current_user.is_superuser and (item.owner_id != current_user.id):
|
||||
raise HTTPException(status_code=403, detail="Not enough permissions")
|
||||
update_dict = item_in.model_dump(exclude_unset=True)
|
||||
item.sqlmodel_update(update_dict)
|
||||
session.add(item)
|
||||
session.commit()
|
||||
session.refresh(item)
|
||||
return item
|
||||
|
||||
|
||||
@router.delete("/{id}")
|
||||
def delete_item(
|
||||
session: SessionDep, current_user: CurrentUser, id: uuid.UUID
|
||||
) -> Message:
|
||||
"""
|
||||
Delete an item.
|
||||
"""
|
||||
item = session.get(Item, id)
|
||||
if not item:
|
||||
raise HTTPException(status_code=404, detail="Item not found")
|
||||
if not current_user.is_superuser and (item.owner_id != current_user.id):
|
||||
raise HTTPException(status_code=403, detail="Not enough permissions")
|
||||
session.delete(item)
|
||||
session.commit()
|
||||
return Message(message="Item deleted successfully")
|
||||
@@ -0,0 +1,123 @@
|
||||
from datetime import timedelta
|
||||
from typing import Annotated, Any
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from fastapi.responses import HTMLResponse
|
||||
from fastapi.security import OAuth2PasswordRequestForm
|
||||
|
||||
from app import crud
|
||||
from app.api.deps import CurrentUser, SessionDep, get_current_active_superuser
|
||||
from app.core import security
|
||||
from app.core.config import settings
|
||||
from app.models import Message, NewPassword, Token, UserPublic, UserUpdate
|
||||
from app.utils import (
|
||||
generate_password_reset_token,
|
||||
generate_reset_password_email,
|
||||
send_email,
|
||||
verify_password_reset_token,
|
||||
)
|
||||
|
||||
router = APIRouter(tags=["login"])
|
||||
|
||||
|
||||
@router.post("/login/access-token")
|
||||
def login_access_token(
|
||||
session: SessionDep, form_data: Annotated[OAuth2PasswordRequestForm, Depends()]
|
||||
) -> Token:
|
||||
"""
|
||||
OAuth2 compatible token login, get an access token for future requests
|
||||
"""
|
||||
user = crud.authenticate(
|
||||
session=session, email=form_data.username, password=form_data.password
|
||||
)
|
||||
if not user:
|
||||
raise HTTPException(status_code=400, detail="Incorrect email or password")
|
||||
elif not user.is_active:
|
||||
raise HTTPException(status_code=400, detail="Inactive user")
|
||||
access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
|
||||
return Token(
|
||||
access_token=security.create_access_token(
|
||||
user.id, expires_delta=access_token_expires
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.post("/login/test-token", response_model=UserPublic)
|
||||
def test_token(current_user: CurrentUser) -> Any:
|
||||
"""
|
||||
Test access token
|
||||
"""
|
||||
return current_user
|
||||
|
||||
|
||||
@router.post("/password-recovery/{email}")
|
||||
def recover_password(email: str, session: SessionDep) -> Message:
|
||||
"""
|
||||
Password Recovery
|
||||
"""
|
||||
user = crud.get_user_by_email(session=session, email=email)
|
||||
|
||||
# Always return the same response to prevent email enumeration attacks
|
||||
# Only send email if user actually exists
|
||||
if user:
|
||||
password_reset_token = generate_password_reset_token(email=email)
|
||||
email_data = generate_reset_password_email(
|
||||
email_to=user.email, email=email, token=password_reset_token
|
||||
)
|
||||
send_email(
|
||||
email_to=user.email,
|
||||
subject=email_data.subject,
|
||||
html_content=email_data.html_content,
|
||||
)
|
||||
return Message(
|
||||
message="If that email is registered, we sent a password recovery link"
|
||||
)
|
||||
|
||||
|
||||
@router.post("/reset-password/")
|
||||
def reset_password(session: SessionDep, body: NewPassword) -> Message:
|
||||
"""
|
||||
Reset password
|
||||
"""
|
||||
email = verify_password_reset_token(token=body.token)
|
||||
if not email:
|
||||
raise HTTPException(status_code=400, detail="Invalid token")
|
||||
user = crud.get_user_by_email(session=session, email=email)
|
||||
if not user:
|
||||
# Don't reveal that the user doesn't exist - use same error as invalid token
|
||||
raise HTTPException(status_code=400, detail="Invalid token")
|
||||
elif not user.is_active:
|
||||
raise HTTPException(status_code=400, detail="Inactive user")
|
||||
user_in_update = UserUpdate(password=body.new_password)
|
||||
crud.update_user(
|
||||
session=session,
|
||||
db_user=user,
|
||||
user_in=user_in_update,
|
||||
)
|
||||
return Message(message="Password updated successfully")
|
||||
|
||||
|
||||
@router.post(
|
||||
"/password-recovery-html-content/{email}",
|
||||
dependencies=[Depends(get_current_active_superuser)],
|
||||
response_class=HTMLResponse,
|
||||
)
|
||||
def recover_password_html_content(email: str, session: SessionDep) -> Any:
|
||||
"""
|
||||
HTML Content for Password Recovery
|
||||
"""
|
||||
user = crud.get_user_by_email(session=session, email=email)
|
||||
|
||||
if not user:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="The user with this username does not exist in the system.",
|
||||
)
|
||||
password_reset_token = generate_password_reset_token(email=email)
|
||||
email_data = generate_reset_password_email(
|
||||
email_to=user.email, email=email, token=password_reset_token
|
||||
)
|
||||
|
||||
return HTMLResponse(
|
||||
content=email_data.html_content, headers={"subject:": email_data.subject}
|
||||
)
|
||||
@@ -0,0 +1,38 @@
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.api.deps import SessionDep
|
||||
from app.core.security import get_password_hash
|
||||
from app.models import (
|
||||
User,
|
||||
UserPublic,
|
||||
)
|
||||
|
||||
router = APIRouter(tags=["private"], prefix="/private")
|
||||
|
||||
|
||||
class PrivateUserCreate(BaseModel):
|
||||
email: str
|
||||
password: str
|
||||
full_name: str
|
||||
is_verified: bool = False
|
||||
|
||||
|
||||
@router.post("/users/", response_model=UserPublic)
|
||||
def create_user(user_in: PrivateUserCreate, session: SessionDep) -> Any:
|
||||
"""
|
||||
Create a new user.
|
||||
"""
|
||||
|
||||
user = User(
|
||||
email=user_in.email,
|
||||
full_name=user_in.full_name,
|
||||
hashed_password=get_password_hash(user_in.password),
|
||||
)
|
||||
|
||||
session.add(user)
|
||||
session.commit()
|
||||
|
||||
return user
|
||||
@@ -0,0 +1,229 @@
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from sqlmodel import col, delete, func, select
|
||||
|
||||
from app import crud
|
||||
from app.api.deps import (
|
||||
CurrentUser,
|
||||
SessionDep,
|
||||
get_current_active_superuser,
|
||||
)
|
||||
from app.core.config import settings
|
||||
from app.core.security import get_password_hash, verify_password
|
||||
from app.models import (
|
||||
Item,
|
||||
Message,
|
||||
UpdatePassword,
|
||||
User,
|
||||
UserCreate,
|
||||
UserPublic,
|
||||
UserRegister,
|
||||
UsersPublic,
|
||||
UserUpdate,
|
||||
UserUpdateMe,
|
||||
)
|
||||
from app.utils import generate_new_account_email, send_email
|
||||
|
||||
router = APIRouter(prefix="/users", tags=["users"])
|
||||
|
||||
|
||||
@router.get(
|
||||
"/",
|
||||
dependencies=[Depends(get_current_active_superuser)],
|
||||
response_model=UsersPublic,
|
||||
)
|
||||
def read_users(session: SessionDep, skip: int = 0, limit: int = 100) -> Any:
|
||||
"""
|
||||
Retrieve users.
|
||||
"""
|
||||
|
||||
count_statement = select(func.count()).select_from(User)
|
||||
count = session.exec(count_statement).one()
|
||||
|
||||
statement = select(User).order_by(User.created_at.desc()).offset(skip).limit(limit)
|
||||
users = session.exec(statement).all()
|
||||
|
||||
return UsersPublic(data=users, count=count)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/", dependencies=[Depends(get_current_active_superuser)], response_model=UserPublic
|
||||
)
|
||||
def create_user(*, session: SessionDep, user_in: UserCreate) -> Any:
|
||||
"""
|
||||
Create new user.
|
||||
"""
|
||||
user = crud.get_user_by_email(session=session, email=user_in.email)
|
||||
if user:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="The user with this email already exists in the system.",
|
||||
)
|
||||
|
||||
user = crud.create_user(session=session, user_create=user_in)
|
||||
if settings.emails_enabled and user_in.email:
|
||||
email_data = generate_new_account_email(
|
||||
email_to=user_in.email, username=user_in.email, password=user_in.password
|
||||
)
|
||||
send_email(
|
||||
email_to=user_in.email,
|
||||
subject=email_data.subject,
|
||||
html_content=email_data.html_content,
|
||||
)
|
||||
return user
|
||||
|
||||
|
||||
@router.patch("/me", response_model=UserPublic)
|
||||
def update_user_me(
|
||||
*, session: SessionDep, user_in: UserUpdateMe, current_user: CurrentUser
|
||||
) -> Any:
|
||||
"""
|
||||
Update own user.
|
||||
"""
|
||||
|
||||
if user_in.email:
|
||||
existing_user = crud.get_user_by_email(session=session, email=user_in.email)
|
||||
if existing_user and existing_user.id != current_user.id:
|
||||
raise HTTPException(
|
||||
status_code=409, detail="User with this email already exists"
|
||||
)
|
||||
user_data = user_in.model_dump(exclude_unset=True)
|
||||
current_user.sqlmodel_update(user_data)
|
||||
session.add(current_user)
|
||||
session.commit()
|
||||
session.refresh(current_user)
|
||||
return current_user
|
||||
|
||||
|
||||
@router.patch("/me/password", response_model=Message)
|
||||
def update_password_me(
|
||||
*, session: SessionDep, body: UpdatePassword, current_user: CurrentUser
|
||||
) -> Any:
|
||||
"""
|
||||
Update own password.
|
||||
"""
|
||||
verified, _ = verify_password(body.current_password, current_user.hashed_password)
|
||||
if not verified:
|
||||
raise HTTPException(status_code=400, detail="Incorrect password")
|
||||
if body.current_password == body.new_password:
|
||||
raise HTTPException(
|
||||
status_code=400, detail="New password cannot be the same as the current one"
|
||||
)
|
||||
hashed_password = get_password_hash(body.new_password)
|
||||
current_user.hashed_password = hashed_password
|
||||
session.add(current_user)
|
||||
session.commit()
|
||||
return Message(message="Password updated successfully")
|
||||
|
||||
|
||||
@router.get("/me", response_model=UserPublic)
|
||||
def read_user_me(current_user: CurrentUser) -> Any:
|
||||
"""
|
||||
Get current user.
|
||||
"""
|
||||
return current_user
|
||||
|
||||
|
||||
@router.delete("/me", response_model=Message)
|
||||
def delete_user_me(session: SessionDep, current_user: CurrentUser) -> Any:
|
||||
"""
|
||||
Delete own user.
|
||||
"""
|
||||
if current_user.is_superuser:
|
||||
raise HTTPException(
|
||||
status_code=403, detail="Super users are not allowed to delete themselves"
|
||||
)
|
||||
session.delete(current_user)
|
||||
session.commit()
|
||||
return Message(message="User deleted successfully")
|
||||
|
||||
|
||||
@router.post("/signup", response_model=UserPublic)
|
||||
def register_user(session: SessionDep, user_in: UserRegister) -> Any:
|
||||
"""
|
||||
Create new user without the need to be logged in.
|
||||
"""
|
||||
user = crud.get_user_by_email(session=session, email=user_in.email)
|
||||
if user:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="The user with this email already exists in the system",
|
||||
)
|
||||
user_create = UserCreate.model_validate(user_in)
|
||||
user = crud.create_user(session=session, user_create=user_create)
|
||||
return user
|
||||
|
||||
|
||||
@router.get("/{user_id}", response_model=UserPublic)
|
||||
def read_user_by_id(
|
||||
user_id: uuid.UUID, session: SessionDep, current_user: CurrentUser
|
||||
) -> Any:
|
||||
"""
|
||||
Get a specific user by id.
|
||||
"""
|
||||
user = session.get(User, user_id)
|
||||
if user == current_user:
|
||||
return user
|
||||
if not current_user.is_superuser:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="The user doesn't have enough privileges",
|
||||
)
|
||||
if user is None:
|
||||
raise HTTPException(status_code=404, detail="User not found")
|
||||
return user
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/{user_id}",
|
||||
dependencies=[Depends(get_current_active_superuser)],
|
||||
response_model=UserPublic,
|
||||
)
|
||||
def update_user(
|
||||
*,
|
||||
session: SessionDep,
|
||||
user_id: uuid.UUID,
|
||||
user_in: UserUpdate,
|
||||
) -> Any:
|
||||
"""
|
||||
Update a user.
|
||||
"""
|
||||
|
||||
db_user = session.get(User, user_id)
|
||||
if not db_user:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="The user with this id does not exist in the system",
|
||||
)
|
||||
if user_in.email:
|
||||
existing_user = crud.get_user_by_email(session=session, email=user_in.email)
|
||||
if existing_user and existing_user.id != user_id:
|
||||
raise HTTPException(
|
||||
status_code=409, detail="User with this email already exists"
|
||||
)
|
||||
|
||||
db_user = crud.update_user(session=session, db_user=db_user, user_in=user_in)
|
||||
return db_user
|
||||
|
||||
|
||||
@router.delete("/{user_id}", dependencies=[Depends(get_current_active_superuser)])
|
||||
def delete_user(
|
||||
session: SessionDep, current_user: CurrentUser, user_id: uuid.UUID
|
||||
) -> Message:
|
||||
"""
|
||||
Delete a user.
|
||||
"""
|
||||
user = session.get(User, user_id)
|
||||
if not user:
|
||||
raise HTTPException(status_code=404, detail="User not found")
|
||||
if user == current_user:
|
||||
raise HTTPException(
|
||||
status_code=403, detail="Super users are not allowed to delete themselves"
|
||||
)
|
||||
statement = delete(Item).where(col(Item.owner_id) == user_id)
|
||||
session.exec(statement) # type: ignore
|
||||
session.delete(user)
|
||||
session.commit()
|
||||
return Message(message="User deleted successfully")
|
||||
@@ -0,0 +1,31 @@
|
||||
from fastapi import APIRouter, Depends
|
||||
from pydantic.networks import EmailStr
|
||||
|
||||
from app.api.deps import get_current_active_superuser
|
||||
from app.models import Message
|
||||
from app.utils import generate_test_email, send_email
|
||||
|
||||
router = APIRouter(prefix="/utils", tags=["utils"])
|
||||
|
||||
|
||||
@router.post(
|
||||
"/test-email/",
|
||||
dependencies=[Depends(get_current_active_superuser)],
|
||||
status_code=201,
|
||||
)
|
||||
def test_email(email_to: EmailStr) -> Message:
|
||||
"""
|
||||
Test emails.
|
||||
"""
|
||||
email_data = generate_test_email(email_to=email_to)
|
||||
send_email(
|
||||
email_to=email_to,
|
||||
subject=email_data.subject,
|
||||
html_content=email_data.html_content,
|
||||
)
|
||||
return Message(message="Test email sent")
|
||||
|
||||
|
||||
@router.get("/health-check/")
|
||||
async def health_check() -> bool:
|
||||
return True
|
||||
@@ -0,0 +1,39 @@
|
||||
import logging
|
||||
|
||||
from sqlalchemy import Engine
|
||||
from sqlmodel import Session, select
|
||||
from tenacity import after_log, before_log, retry, stop_after_attempt, wait_fixed
|
||||
|
||||
from app.core.db import engine
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
max_tries = 60 * 5 # 5 minutes
|
||||
wait_seconds = 1
|
||||
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(max_tries),
|
||||
wait=wait_fixed(wait_seconds),
|
||||
before=before_log(logger, logging.INFO),
|
||||
after=after_log(logger, logging.WARN),
|
||||
)
|
||||
def init(db_engine: Engine) -> None:
|
||||
try:
|
||||
with Session(db_engine) as session:
|
||||
# Try to create session to check if DB is awake
|
||||
session.exec(select(1))
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
raise e
|
||||
|
||||
|
||||
def main() -> None:
|
||||
logger.info("Initializing service")
|
||||
init(engine)
|
||||
logger.info("Service finished initializing")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,119 @@
|
||||
import secrets
|
||||
import warnings
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
from pydantic import (
|
||||
AnyUrl,
|
||||
BeforeValidator,
|
||||
EmailStr,
|
||||
HttpUrl,
|
||||
PostgresDsn,
|
||||
computed_field,
|
||||
model_validator,
|
||||
)
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
from typing_extensions import Self
|
||||
|
||||
|
||||
def parse_cors(v: Any) -> list[str] | str:
|
||||
if isinstance(v, str) and not v.startswith("["):
|
||||
return [i.strip() for i in v.split(",") if i.strip()]
|
||||
elif isinstance(v, list | str):
|
||||
return v
|
||||
raise ValueError(v)
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
model_config = SettingsConfigDict(
|
||||
# Use top level .env file (one level above ./backend/)
|
||||
env_file="../.env",
|
||||
env_ignore_empty=True,
|
||||
extra="ignore",
|
||||
)
|
||||
API_V1_STR: str = "/api/v1"
|
||||
SECRET_KEY: str = secrets.token_urlsafe(32)
|
||||
# 60 minutes * 24 hours * 8 days = 8 days
|
||||
ACCESS_TOKEN_EXPIRE_MINUTES: int = 60 * 24 * 8
|
||||
FRONTEND_HOST: str = "http://localhost:5173"
|
||||
ENVIRONMENT: Literal["local", "staging", "production"] = "local"
|
||||
|
||||
BACKEND_CORS_ORIGINS: Annotated[
|
||||
list[AnyUrl] | str, BeforeValidator(parse_cors)
|
||||
] = []
|
||||
|
||||
@computed_field # type: ignore[prop-decorator]
|
||||
@property
|
||||
def all_cors_origins(self) -> list[str]:
|
||||
return [str(origin).rstrip("/") for origin in self.BACKEND_CORS_ORIGINS] + [
|
||||
self.FRONTEND_HOST
|
||||
]
|
||||
|
||||
PROJECT_NAME: str
|
||||
SENTRY_DSN: HttpUrl | None = None
|
||||
POSTGRES_SERVER: str
|
||||
POSTGRES_PORT: int = 5432
|
||||
POSTGRES_USER: str
|
||||
POSTGRES_PASSWORD: str = ""
|
||||
POSTGRES_DB: str = ""
|
||||
|
||||
@computed_field # type: ignore[prop-decorator]
|
||||
@property
|
||||
def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn:
|
||||
return PostgresDsn.build(
|
||||
scheme="postgresql+psycopg",
|
||||
username=self.POSTGRES_USER,
|
||||
password=self.POSTGRES_PASSWORD,
|
||||
host=self.POSTGRES_SERVER,
|
||||
port=self.POSTGRES_PORT,
|
||||
path=self.POSTGRES_DB,
|
||||
)
|
||||
|
||||
SMTP_TLS: bool = True
|
||||
SMTP_SSL: bool = False
|
||||
SMTP_PORT: int = 587
|
||||
SMTP_HOST: str | None = None
|
||||
SMTP_USER: str | None = None
|
||||
SMTP_PASSWORD: str | None = None
|
||||
EMAILS_FROM_EMAIL: EmailStr | None = None
|
||||
EMAILS_FROM_NAME: str | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _set_default_emails_from(self) -> Self:
|
||||
if not self.EMAILS_FROM_NAME:
|
||||
self.EMAILS_FROM_NAME = self.PROJECT_NAME
|
||||
return self
|
||||
|
||||
EMAIL_RESET_TOKEN_EXPIRE_HOURS: int = 48
|
||||
|
||||
@computed_field # type: ignore[prop-decorator]
|
||||
@property
|
||||
def emails_enabled(self) -> bool:
|
||||
return bool(self.SMTP_HOST and self.EMAILS_FROM_EMAIL)
|
||||
|
||||
EMAIL_TEST_USER: EmailStr = "test@example.com"
|
||||
FIRST_SUPERUSER: EmailStr
|
||||
FIRST_SUPERUSER_PASSWORD: str
|
||||
|
||||
def _check_default_secret(self, var_name: str, value: str | None) -> None:
|
||||
if value == "changethis":
|
||||
message = (
|
||||
f'The value of {var_name} is "changethis", '
|
||||
"for security, please change it, at least for deployments."
|
||||
)
|
||||
if self.ENVIRONMENT == "local":
|
||||
warnings.warn(message, stacklevel=1)
|
||||
else:
|
||||
raise ValueError(message)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _enforce_non_default_secrets(self) -> Self:
|
||||
self._check_default_secret("SECRET_KEY", self.SECRET_KEY)
|
||||
self._check_default_secret("POSTGRES_PASSWORD", self.POSTGRES_PASSWORD)
|
||||
self._check_default_secret(
|
||||
"FIRST_SUPERUSER_PASSWORD", self.FIRST_SUPERUSER_PASSWORD
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
settings = Settings() # type: ignore
|
||||
@@ -0,0 +1,33 @@
|
||||
from sqlmodel import Session, create_engine, select
|
||||
|
||||
from app import crud
|
||||
from app.core.config import settings
|
||||
from app.models import User, UserCreate
|
||||
|
||||
engine = create_engine(str(settings.SQLALCHEMY_DATABASE_URI))
|
||||
|
||||
|
||||
# make sure all SQLModel models are imported (app.models) before initializing DB
|
||||
# otherwise, SQLModel might fail to initialize relationships properly
|
||||
# for more details: https://github.com/fastapi/full-stack-fastapi-template/issues/28
|
||||
|
||||
|
||||
def init_db(session: Session) -> None:
|
||||
# Tables should be created with Alembic migrations
|
||||
# But if you don't want to use migrations, create
|
||||
# the tables un-commenting the next lines
|
||||
# from sqlmodel import SQLModel
|
||||
|
||||
# This works because the models are already imported and registered from app.models
|
||||
# SQLModel.metadata.create_all(engine)
|
||||
|
||||
user = session.exec(
|
||||
select(User).where(User.email == settings.FIRST_SUPERUSER)
|
||||
).first()
|
||||
if not user:
|
||||
user_in = UserCreate(
|
||||
email=settings.FIRST_SUPERUSER,
|
||||
password=settings.FIRST_SUPERUSER_PASSWORD,
|
||||
is_superuser=True,
|
||||
)
|
||||
user = crud.create_user(session=session, user_create=user_in)
|
||||
@@ -0,0 +1,36 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
|
||||
import jwt
|
||||
from pwdlib import PasswordHash
|
||||
from pwdlib.hashers.argon2 import Argon2Hasher
|
||||
from pwdlib.hashers.bcrypt import BcryptHasher
|
||||
|
||||
from app.core.config import settings
|
||||
|
||||
password_hash = PasswordHash(
|
||||
(
|
||||
Argon2Hasher(),
|
||||
BcryptHasher(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
ALGORITHM = "HS256"
|
||||
|
||||
|
||||
def create_access_token(subject: str | Any, expires_delta: timedelta) -> str:
|
||||
expire = datetime.now(timezone.utc) + expires_delta
|
||||
to_encode = {"exp": expire, "sub": str(subject)}
|
||||
encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=ALGORITHM)
|
||||
return encoded_jwt
|
||||
|
||||
|
||||
def verify_password(
|
||||
plain_password: str, hashed_password: str
|
||||
) -> tuple[bool, str | None]:
|
||||
return password_hash.verify_and_update(plain_password, hashed_password)
|
||||
|
||||
|
||||
def get_password_hash(password: str) -> str:
|
||||
return password_hash.hash(password)
|
||||
@@ -0,0 +1,68 @@
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from sqlmodel import Session, select
|
||||
|
||||
from app.core.security import get_password_hash, verify_password
|
||||
from app.models import Item, ItemCreate, User, UserCreate, UserUpdate
|
||||
|
||||
|
||||
def create_user(*, session: Session, user_create: UserCreate) -> User:
|
||||
db_obj = User.model_validate(
|
||||
user_create, update={"hashed_password": get_password_hash(user_create.password)}
|
||||
)
|
||||
session.add(db_obj)
|
||||
session.commit()
|
||||
session.refresh(db_obj)
|
||||
return db_obj
|
||||
|
||||
|
||||
def update_user(*, session: Session, db_user: User, user_in: UserUpdate) -> Any:
|
||||
user_data = user_in.model_dump(exclude_unset=True)
|
||||
extra_data = {}
|
||||
if "password" in user_data:
|
||||
password = user_data["password"]
|
||||
hashed_password = get_password_hash(password)
|
||||
extra_data["hashed_password"] = hashed_password
|
||||
db_user.sqlmodel_update(user_data, update=extra_data)
|
||||
session.add(db_user)
|
||||
session.commit()
|
||||
session.refresh(db_user)
|
||||
return db_user
|
||||
|
||||
|
||||
def get_user_by_email(*, session: Session, email: str) -> User | None:
|
||||
statement = select(User).where(User.email == email)
|
||||
session_user = session.exec(statement).first()
|
||||
return session_user
|
||||
|
||||
|
||||
# Dummy hash to use for timing attack prevention when user is not found
|
||||
# This is an Argon2 hash of a random password, used to ensure constant-time comparison
|
||||
DUMMY_HASH = "$argon2id$v=19$m=65536,t=3,p=4$MjQyZWE1MzBjYjJlZTI0Yw$YTU4NGM5ZTZmYjE2NzZlZjY0ZWY3ZGRkY2U2OWFjNjk"
|
||||
|
||||
|
||||
def authenticate(*, session: Session, email: str, password: str) -> User | None:
|
||||
db_user = get_user_by_email(session=session, email=email)
|
||||
if not db_user:
|
||||
# Prevent timing attacks by running password verification even when user doesn't exist
|
||||
# This ensures the response time is similar whether or not the email exists
|
||||
verify_password(password, DUMMY_HASH)
|
||||
return None
|
||||
verified, updated_password_hash = verify_password(password, db_user.hashed_password)
|
||||
if not verified:
|
||||
return None
|
||||
if updated_password_hash:
|
||||
db_user.hashed_password = updated_password_hash
|
||||
session.add(db_user)
|
||||
session.commit()
|
||||
session.refresh(db_user)
|
||||
return db_user
|
||||
|
||||
|
||||
def create_item(*, session: Session, item_in: ItemCreate, owner_id: uuid.UUID) -> Item:
|
||||
db_item = Item.model_validate(item_in, update={"owner_id": owner_id})
|
||||
session.add(db_item)
|
||||
session.commit()
|
||||
session.refresh(db_item)
|
||||
return db_item
|
||||
@@ -0,0 +1,25 @@
|
||||
<!doctype html><html xmlns="http://www.w3.org/1999/xhtml" xmlns:v="urn:schemas-microsoft-com:vml" xmlns:o="urn:schemas-microsoft-com:office:office"><head><title></title><!--[if !mso]><!-- --><meta http-equiv="X-UA-Compatible" content="IE=edge"><!--<![endif]--><meta http-equiv="Content-Type" content="text/html; charset=UTF-8"><meta name="viewport" content="width=device-width,initial-scale=1"><style type="text/css">#outlook a { padding:0; }
|
||||
.ReadMsgBody { width:100%; }
|
||||
.ExternalClass { width:100%; }
|
||||
.ExternalClass * { line-height:100%; }
|
||||
body { margin:0;padding:0;-webkit-text-size-adjust:100%;-ms-text-size-adjust:100%; }
|
||||
table, td { border-collapse:collapse;mso-table-lspace:0pt;mso-table-rspace:0pt; }
|
||||
img { border:0;height:auto;line-height:100%; outline:none;text-decoration:none;-ms-interpolation-mode:bicubic; }
|
||||
p { display:block;margin:13px 0; }</style><!--[if !mso]><!--><style type="text/css">@media only screen and (max-width:480px) {
|
||||
@-ms-viewport { width:320px; }
|
||||
@viewport { width:320px; }
|
||||
}</style><!--<![endif]--><!--[if mso]>
|
||||
<xml>
|
||||
<o:OfficeDocumentSettings>
|
||||
<o:AllowPNG/>
|
||||
<o:PixelsPerInch>96</o:PixelsPerInch>
|
||||
</o:OfficeDocumentSettings>
|
||||
</xml>
|
||||
<![endif]--><!--[if lte mso 11]>
|
||||
<style type="text/css">
|
||||
.outlook-group-fix { width:100% !important; }
|
||||
</style>
|
||||
<![endif]--><!--[if !mso]><!--><link href="https://fonts.googleapis.com/css?family=Ubuntu:300,400,500,700" rel="stylesheet" type="text/css"><style type="text/css">@import url(https://fonts.googleapis.com/css?family=Ubuntu:300,400,500,700);</style><!--<![endif]--><style type="text/css">@media only screen and (min-width:480px) {
|
||||
.mj-column-per-100 { width:100% !important; max-width: 100%; }
|
||||
}</style><style type="text/css"></style></head><body style="background-color:#fafbfc;"><div style="background-color:#fafbfc;"><!--[if mso | IE]><table align="center" border="0" cellpadding="0" cellspacing="0" class="" style="width:600px;" width="600" ><tr><td style="line-height:0px;font-size:0px;mso-line-height-rule:exactly;"><![endif]--><div style="background:#ffffff;background-color:#ffffff;Margin:0px auto;max-width:600px;"><table align="center" border="0" cellpadding="0" cellspacing="0" role="presentation" style="background:#ffffff;background-color:#ffffff;width:100%;"><tbody><tr><td style="direction:ltr;font-size:0px;padding:40px 20px;text-align:center;vertical-align:top;"><!--[if mso | IE]><table role="presentation" border="0" cellpadding="0" cellspacing="0"><tr><td class="" style="vertical-align:middle;width:560px;" ><![endif]--><div class="mj-column-per-100 outlook-group-fix" style="font-size:13px;text-align:left;direction:ltr;display:inline-block;vertical-align:middle;width:100%;"><table border="0" cellpadding="0" cellspacing="0" role="presentation" style="vertical-align:middle;" width="100%"><tr><td align="center" style="font-size:0px;padding:35px;word-break:break-word;"><div style="font-family:Ubuntu, Helvetica, Arial, sans-serif;font-size:20px;line-height:1;text-align:center;color:#333333;">{{ project_name }} - New Account</div></td></tr><tr><td align="center" style="font-size:0px;padding:10px 25px;padding-right:25px;padding-left:25px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:16px;line-height:1;text-align:center;color:#555555;"><span>Welcome to your new account!</span></div></td></tr><tr><td align="center" style="font-size:0px;padding:10px 25px;padding-right:25px;padding-left:25px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:16px;line-height:1;text-align:center;color:#555555;">Here are your account details:</div></td></tr><tr><td align="center" style="font-size:0px;padding:10px 25px;padding-right:25px;padding-left:25px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:16px;line-height:1;text-align:center;color:#555555;">Username: {{ username }}</div></td></tr><tr><td align="center" style="font-size:0px;padding:10px 25px;padding-right:25px;padding-left:25px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:16px;line-height:1;text-align:center;color:#555555;">Password: {{ password }}</div></td></tr><tr><td align="center" vertical-align="middle" style="font-size:0px;padding:15px 30px;word-break:break-word;"><table border="0" cellpadding="0" cellspacing="0" role="presentation" style="border-collapse:separate;line-height:100%;"><tr><td align="center" bgcolor="#009688" role="presentation" style="border:none;border-radius:8px;cursor:auto;padding:10px 25px;background:#009688;" valign="middle"><a href="{{ link }}" style="background:#009688;color:#ffffff;font-family:Ubuntu, Helvetica, Arial, sans-serif;font-size:18px;font-weight:normal;line-height:120%;Margin:0;text-decoration:none;text-transform:none;" target="_blank">Go to Dashboard</a></td></tr></table></td></tr><tr><td style="font-size:0px;padding:10px 25px;word-break:break-word;"><p style="border-top:solid 2px #cccccc;font-size:1;margin:0px auto;width:100%;"></p><!--[if mso | IE]><table align="center" border="0" cellpadding="0" cellspacing="0" style="border-top:solid 2px #cccccc;font-size:1;margin:0px auto;width:510px;" role="presentation" width="510px" ><tr><td style="height:0;line-height:0;">
|
||||
</td></tr></table><![endif]--></td></tr></table></div><!--[if mso | IE]></td></tr></table><![endif]--></td></tr></tbody></table></div><!--[if mso | IE]></td></tr></table><![endif]--></div></body></html>
|
||||
@@ -0,0 +1,25 @@
|
||||
<!doctype html><html xmlns="http://www.w3.org/1999/xhtml" xmlns:v="urn:schemas-microsoft-com:vml" xmlns:o="urn:schemas-microsoft-com:office:office"><head><title></title><!--[if !mso]><!-- --><meta http-equiv="X-UA-Compatible" content="IE=edge"><!--<![endif]--><meta http-equiv="Content-Type" content="text/html; charset=UTF-8"><meta name="viewport" content="width=device-width,initial-scale=1"><style type="text/css">#outlook a { padding:0; }
|
||||
.ReadMsgBody { width:100%; }
|
||||
.ExternalClass { width:100%; }
|
||||
.ExternalClass * { line-height:100%; }
|
||||
body { margin:0;padding:0;-webkit-text-size-adjust:100%;-ms-text-size-adjust:100%; }
|
||||
table, td { border-collapse:collapse;mso-table-lspace:0pt;mso-table-rspace:0pt; }
|
||||
img { border:0;height:auto;line-height:100%; outline:none;text-decoration:none;-ms-interpolation-mode:bicubic; }
|
||||
p { display:block;margin:13px 0; }</style><!--[if !mso]><!--><style type="text/css">@media only screen and (max-width:480px) {
|
||||
@-ms-viewport { width:320px; }
|
||||
@viewport { width:320px; }
|
||||
}</style><!--<![endif]--><!--[if mso]>
|
||||
<xml>
|
||||
<o:OfficeDocumentSettings>
|
||||
<o:AllowPNG/>
|
||||
<o:PixelsPerInch>96</o:PixelsPerInch>
|
||||
</o:OfficeDocumentSettings>
|
||||
</xml>
|
||||
<![endif]--><!--[if lte mso 11]>
|
||||
<style type="text/css">
|
||||
.outlook-group-fix { width:100% !important; }
|
||||
</style>
|
||||
<![endif]--><!--[if !mso]><!--><link href="https://fonts.googleapis.com/css?family=Ubuntu:300,400,500,700" rel="stylesheet" type="text/css"><style type="text/css">@import url(https://fonts.googleapis.com/css?family=Ubuntu:300,400,500,700);</style><!--<![endif]--><style type="text/css">@media only screen and (min-width:480px) {
|
||||
.mj-column-per-100 { width:100% !important; max-width: 100%; }
|
||||
}</style><style type="text/css"></style></head><body style="background-color:#fafbfc;"><div style="background-color:#fafbfc;"><!--[if mso | IE]><table align="center" border="0" cellpadding="0" cellspacing="0" class="" style="width:600px;" width="600" ><tr><td style="line-height:0px;font-size:0px;mso-line-height-rule:exactly;"><![endif]--><div style="background:#ffffff;background-color:#ffffff;Margin:0px auto;max-width:600px;"><table align="center" border="0" cellpadding="0" cellspacing="0" role="presentation" style="background:#ffffff;background-color:#ffffff;width:100%;"><tbody><tr><td style="direction:ltr;font-size:0px;padding:40px 20px;text-align:center;vertical-align:top;"><!--[if mso | IE]><table role="presentation" border="0" cellpadding="0" cellspacing="0"><tr><td class="" style="vertical-align:middle;width:560px;" ><![endif]--><div class="mj-column-per-100 outlook-group-fix" style="font-size:13px;text-align:left;direction:ltr;display:inline-block;vertical-align:middle;width:100%;"><table border="0" cellpadding="0" cellspacing="0" role="presentation" style="vertical-align:middle;" width="100%"><tr><td align="center" style="font-size:0px;padding:35px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:20px;line-height:1;text-align:center;color:#333333;">{{ project_name }} - Password Recovery</div></td></tr><tr><td align="center" style="font-size:0px;padding:10px 25px;padding-right:25px;padding-left:25px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:16px;line-height:1;text-align:center;color:#555555;"><span>Hello {{ username }}</span></div></td></tr><tr><td align="center" style="font-size:0px;padding:10px 25px;padding-right:25px;padding-left:25px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:16px;line-height:1;text-align:center;color:#555555;">We've received a request to reset your password. You can do it by clicking the button below:</div></td></tr><tr><td align="center" vertical-align="middle" style="font-size:0px;padding:15px 30px;word-break:break-word;"><table border="0" cellpadding="0" cellspacing="0" role="presentation" style="border-collapse:separate;line-height:100%;"><tr><td align="center" bgcolor="#009688" role="presentation" style="border:none;border-radius:8px;cursor:auto;padding:10px 25px;background:#009688;" valign="middle"><a href="{{ link }}" style="background:#009688;color:#ffffff;font-family:Ubuntu, Helvetica, Arial, sans-serif;font-size:18px;font-weight:normal;line-height:120%;Margin:0;text-decoration:none;text-transform:none;" target="_blank">Reset password</a></td></tr></table></td></tr><tr><td align="center" style="font-size:0px;padding:10px 25px;padding-right:25px;padding-left:25px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:16px;line-height:1;text-align:center;color:#555555;">Or copy and paste the following link into your browser:</div></td></tr><tr><td align="center" style="font-size:0px;padding:10px 25px;padding-right:25px;padding-left:25px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:16px;line-height:1;text-align:center;color:#555555;"><a href="{{ link }}">{{ link }}</a></div></td></tr><tr><td align="center" style="font-size:0px;padding:10px 25px;padding-right:25px;padding-left:25px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:16px;line-height:1;text-align:center;color:#555555;">This password will expire in {{ valid_hours }} hours.</div></td></tr><tr><td style="font-size:0px;padding:10px 25px;word-break:break-word;"><p style="border-top:solid 2px #cccccc;font-size:1;margin:0px auto;width:100%;"></p><!--[if mso | IE]><table align="center" border="0" cellpadding="0" cellspacing="0" style="border-top:solid 2px #cccccc;font-size:1;margin:0px auto;width:510px;" role="presentation" width="510px" ><tr><td style="height:0;line-height:0;">
|
||||
</td></tr></table><![endif]--></td></tr><tr><td align="center" style="font-size:0px;padding:10px 25px;padding-right:25px;padding-left:25px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:14px;line-height:1;text-align:center;color:#555555;">If you didn't request a password recovery you can disregard this email.</div></td></tr></table></div><!--[if mso | IE]></td></tr></table><![endif]--></td></tr></tbody></table></div><!--[if mso | IE]></td></tr></table><![endif]--></div></body></html>
|
||||
@@ -0,0 +1,25 @@
|
||||
<!doctype html><html xmlns="http://www.w3.org/1999/xhtml" xmlns:v="urn:schemas-microsoft-com:vml" xmlns:o="urn:schemas-microsoft-com:office:office"><head><title></title><!--[if !mso]><!-- --><meta http-equiv="X-UA-Compatible" content="IE=edge"><!--<![endif]--><meta http-equiv="Content-Type" content="text/html; charset=UTF-8"><meta name="viewport" content="width=device-width,initial-scale=1"><style type="text/css">#outlook a { padding:0; }
|
||||
.ReadMsgBody { width:100%; }
|
||||
.ExternalClass { width:100%; }
|
||||
.ExternalClass * { line-height:100%; }
|
||||
body { margin:0;padding:0;-webkit-text-size-adjust:100%;-ms-text-size-adjust:100%; }
|
||||
table, td { border-collapse:collapse;mso-table-lspace:0pt;mso-table-rspace:0pt; }
|
||||
img { border:0;height:auto;line-height:100%; outline:none;text-decoration:none;-ms-interpolation-mode:bicubic; }
|
||||
p { display:block;margin:13px 0; }</style><!--[if !mso]><!--><style type="text/css">@media only screen and (max-width:480px) {
|
||||
@-ms-viewport { width:320px; }
|
||||
@viewport { width:320px; }
|
||||
}</style><!--<![endif]--><!--[if mso]>
|
||||
<xml>
|
||||
<o:OfficeDocumentSettings>
|
||||
<o:AllowPNG/>
|
||||
<o:PixelsPerInch>96</o:PixelsPerInch>
|
||||
</o:OfficeDocumentSettings>
|
||||
</xml>
|
||||
<![endif]--><!--[if lte mso 11]>
|
||||
<style type="text/css">
|
||||
.outlook-group-fix { width:100% !important; }
|
||||
</style>
|
||||
<![endif]--><style type="text/css">@media only screen and (min-width:480px) {
|
||||
.mj-column-per-100 { width:100% !important; max-width: 100%; }
|
||||
}</style><style type="text/css"></style></head><body style="background-color:#fafbfc;"><div style="background-color:#fafbfc;"><!--[if mso | IE]><table align="center" border="0" cellpadding="0" cellspacing="0" class="" style="width:600px;" width="600" ><tr><td style="line-height:0px;font-size:0px;mso-line-height-rule:exactly;"><![endif]--><div style="background:#ffffff;background-color:#ffffff;Margin:0px auto;max-width:600px;"><table align="center" border="0" cellpadding="0" cellspacing="0" role="presentation" style="background:#ffffff;background-color:#ffffff;width:100%;"><tbody><tr><td style="direction:ltr;font-size:0px;padding:40px 20px;text-align:center;vertical-align:top;"><!--[if mso | IE]><table role="presentation" border="0" cellpadding="0" cellspacing="0"><tr><td class="" style="vertical-align:middle;width:560px;" ><![endif]--><div class="mj-column-per-100 outlook-group-fix" style="font-size:13px;text-align:left;direction:ltr;display:inline-block;vertical-align:middle;width:100%;"><table border="0" cellpadding="0" cellspacing="0" role="presentation" style="vertical-align:middle;" width="100%"><tr><td align="center" style="font-size:0px;padding:35px;word-break:break-word;"><div style="font-family:Arial, Helvetica, sans-serif;font-size:20px;line-height:1;text-align:center;color:#333333;">{{ project_name }}</div></td></tr><tr><td align="center" style="font-size:0px;padding:10px 25px;padding-right:25px;padding-left:25px;word-break:break-word;"><div style="font-family:, sans-serif;font-size:16px;line-height:1;text-align:center;color:#555555;"><span>Test email for: {{ email }}</span></div></td></tr><tr><td style="font-size:0px;padding:10px 25px;word-break:break-word;"><p style="border-top:solid 2px #cccccc;font-size:1;margin:0px auto;width:100%;"></p><!--[if mso | IE]><table align="center" border="0" cellpadding="0" cellspacing="0" style="border-top:solid 2px #cccccc;font-size:1;margin:0px auto;width:510px;" role="presentation" width="510px" ><tr><td style="height:0;line-height:0;">
|
||||
</td></tr></table><![endif]--></td></tr></table></div><!--[if mso | IE]></td></tr></table><![endif]--></td></tr></tbody></table></div><!--[if mso | IE]></td></tr></table><![endif]--></div></body></html>
|
||||
@@ -0,0 +1,15 @@
|
||||
<mjml>
|
||||
<mj-body background-color="#fafbfc">
|
||||
<mj-section background-color="#fff" padding="40px 20px">
|
||||
<mj-column vertical-align="middle" width="100%">
|
||||
<mj-text align="center" padding="35px" font-size="20px" color="#333">{{ project_name }} - New Account</mj-text>
|
||||
<mj-text align="center" font-size="16px" padding-left="25px" padding-right="25px" font-family="Arial, Helvetica, sans-serif" color="#555"><span>Welcome to your new account!</span></mj-text>
|
||||
<mj-text align="center" font-size="16px" padding-left="25px" padding-right="25px" font-family="Arial, Helvetica, sans-serif" color="#555">Here are your account details:</mj-text>
|
||||
<mj-text align="center" font-size="16px" padding-left="25px" padding-right="25px" font-family="Arial, Helvetica, sans-serif" color="#555">Username: {{ username }}</mj-text>
|
||||
<mj-text align="center" font-size="16px" padding-left="25px" padding-right="25px" font-family="Arial, Helvetica, sans-serif" color="#555">Password: {{ password }}</mj-text>
|
||||
<mj-button align="center" font-size="18px" background-color="#009688" border-radius="8px" color="#fff" href="{{ link }}" padding="15px 30px">Go to Dashboard</mj-button>
|
||||
<mj-divider border-color="#ccc" border-width="2px"></mj-divider>
|
||||
</mj-column>
|
||||
</mj-section>
|
||||
</mj-body>
|
||||
</mjml>
|
||||
@@ -0,0 +1,17 @@
|
||||
<mjml>
|
||||
<mj-body background-color="#fafbfc">
|
||||
<mj-section background-color="#fff" padding="40px 20px">
|
||||
<mj-column vertical-align="middle" width="100%">
|
||||
<mj-text align="center" padding="35px" font-size="20px" font-family="Arial, Helvetica, sans-serif" color="#333">{{ project_name }} - Password Recovery</mj-text>
|
||||
<mj-text align="center" font-size="16px" padding-left="25px" padding-right="25px" font-family="Arial, Helvetica, sans-serif" color="#555"><span>Hello {{ username }}</span></mj-text>
|
||||
<mj-text align="center" font-size="16px" padding-left="25px" padding-right="25px" font-family="Arial, Helvetica, sans-serif" color="#555">We've received a request to reset your password. You can do it by clicking the button below:</mj-text>
|
||||
<mj-button align="center" font-size="18px" background-color="#009688" border-radius="8px" color="#fff" href="{{ link }}" padding="15px 30px">Reset password</mj-button>
|
||||
<mj-text align="center" font-size="16px" padding-left="25px" padding-right="25px" font-family="Arial, Helvetica, sans-serif" color="#555">Or copy and paste the following link into your browser:</mj-text>
|
||||
<mj-text align="center" font-size="16px" padding-left="25px" padding-right="25px" font-family="Arial, Helvetica, sans-serif" color="#555"><a href="{{ link }}">{{ link }}</a></mj-text>
|
||||
<mj-text align="center" font-size="16px" padding-left="25px" padding-right="25px" font-family="Arial, Helvetica, sans-serif" color="#555">This password will expire in {{ valid_hours }} hours.</mj-text>
|
||||
<mj-divider border-color="#ccc" border-width="2px"></mj-divider>
|
||||
<mj-text align="center" font-size="14px" padding-left="25px" padding-right="25px" font-family="Arial, Helvetica, sans-serif" color="#555">If you didn't request a password recovery you can disregard this email.</mj-text>
|
||||
</mj-column>
|
||||
</mj-section>
|
||||
</mj-body>
|
||||
</mjml>
|
||||
@@ -0,0 +1,11 @@
|
||||
<mjml>
|
||||
<mj-body background-color="#fafbfc">
|
||||
<mj-section background-color="#fff" padding="40px 20px">
|
||||
<mj-column vertical-align="middle" width="100%">
|
||||
<mj-text align="center" padding="35px" font-size="20px" font-family="Arial, Helvetica, sans-serif" color="#333">{{ project_name }}</mj-text>
|
||||
<mj-text align="center" font-size="16px" padding-left="25px" padding-right="25px" font-family=", sans-serif" color="#555"><span>Test email for: {{ email }}</span></mj-text>
|
||||
<mj-divider border-color="#ccc" border-width="2px"></mj-divider>
|
||||
</mj-column>
|
||||
</mj-section>
|
||||
</mj-body>
|
||||
</mjml>
|
||||
@@ -0,0 +1,813 @@
|
||||
"""
|
||||
Pipeline Controller with file watching for dynamic node updates.
|
||||
|
||||
This module provides a PipelineController class that manages a pipeline,
|
||||
watches for changes in node files, and dynamically updates the graph.
|
||||
Handles errors gracefully to keep the pipeline intact when nodes have issues.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import importlib.util
|
||||
import logging
|
||||
import sys
|
||||
import traceback
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, TYPE_CHECKING
|
||||
from watchfiles import awatch, Change
|
||||
|
||||
from nodes import Node, HttpNode, MqttNode, InfluxDbNode, DelayNode
|
||||
from pipeline import Pipeline
|
||||
from state import StateBackend, MemoryState, RedisState
|
||||
from util import Message
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import FastAPI
|
||||
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NodeStatus(Enum):
|
||||
"""Status of a loaded node module."""
|
||||
|
||||
ACTIVE = "active"
|
||||
ERROR = "error"
|
||||
DISABLED = "disabled"
|
||||
|
||||
|
||||
@dataclass
|
||||
class LoadedNode:
|
||||
"""Represents a loaded node with metadata."""
|
||||
|
||||
node: Node | None
|
||||
file_path: Path
|
||||
file_hash: str
|
||||
status: NodeStatus
|
||||
error_message: str | None = None
|
||||
last_modified: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class PipelineState:
|
||||
"""Holds the current state of the pipeline controller."""
|
||||
|
||||
nodes: dict[str, LoadedNode] = field(default_factory=dict)
|
||||
pipeline: Pipeline | None = None
|
||||
is_running: bool = False
|
||||
|
||||
|
||||
class PipelineController:
|
||||
"""
|
||||
Controller for managing a dynamic pipeline with file watching.
|
||||
|
||||
Watches a directory for node files, dynamically loads/unloads nodes,
|
||||
and rebuilds the pipeline graph when changes are detected. Errors in
|
||||
individual nodes don't affect the rest of the pipeline.
|
||||
|
||||
:param nodes_dir: Directory containing node files.
|
||||
:type nodes_dir: str | Path
|
||||
:param state_backend: State backend for the pipeline.
|
||||
:type state_backend: StateBackend | None
|
||||
:param max_workers: Maximum workers for parallel execution.
|
||||
:type max_workers: int
|
||||
:param node_params: Default parameters passed to nodes.
|
||||
:type node_params: dict | None
|
||||
|
||||
:example:
|
||||
>>> controller = PipelineController(
|
||||
... nodes_dir="./nodes",
|
||||
... state_backend=RedisState(host="localhost"),
|
||||
... max_workers=4,
|
||||
... )
|
||||
>>> await controller.start()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
nodes_dir: str | Path,
|
||||
state_backend: StateBackend | None = None,
|
||||
max_workers: int = 4,
|
||||
node_params: dict | None = None,
|
||||
fastapi_app: "FastAPI | None" = None,
|
||||
) -> None:
|
||||
self.nodes_dir = Path(nodes_dir)
|
||||
self.state_backend = state_backend or MemoryState()
|
||||
self.max_workers = max_workers
|
||||
self.node_params = node_params or {}
|
||||
self._fastapi_app: "FastAPI | None" = fastapi_app
|
||||
|
||||
self._state = PipelineState()
|
||||
self._watch_task: asyncio.Task | None = None
|
||||
self._rebuild_lock = asyncio.Lock()
|
||||
self._shutdown_event = asyncio.Event()
|
||||
|
||||
# Callbacks for external integrations (e.g., frontend notifications)
|
||||
self._on_node_loaded: list[Callable[[str, Node], None]] = []
|
||||
self._on_node_error: list[Callable[[str, str], None]] = []
|
||||
self._on_node_removed: list[Callable[[str], None]] = []
|
||||
self._on_pipeline_rebuilt: list[Callable[[Pipeline], None]] = []
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Event Callbacks
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def on_node_loaded(self, callback: Callable[[str, Node], None]) -> None:
|
||||
"""Register a callback for when a node is successfully loaded."""
|
||||
self._on_node_loaded.append(callback)
|
||||
|
||||
def on_node_error(self, callback: Callable[[str, str], None]) -> None:
|
||||
"""Register a callback for when a node fails to load."""
|
||||
self._on_node_error.append(callback)
|
||||
|
||||
def on_node_removed(self, callback: Callable[[str], None]) -> None:
|
||||
"""Register a callback for when a node is removed."""
|
||||
self._on_node_removed.append(callback)
|
||||
|
||||
def on_pipeline_rebuilt(self, callback: Callable[[Pipeline], None]) -> None:
|
||||
"""Register a callback for when the pipeline is rebuilt."""
|
||||
self._on_pipeline_rebuilt.append(callback)
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Node Loading
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def _compute_file_hash(self, file_path: Path) -> str:
|
||||
"""Compute MD5 hash of file contents for change detection."""
|
||||
content = file_path.read_bytes()
|
||||
return hashlib.md5(content).hexdigest()
|
||||
|
||||
def _load_node_from_file(self, file_path: Path) -> LoadedNode:
|
||||
"""
|
||||
Load a node from a Python file.
|
||||
|
||||
The file should define either:
|
||||
- A `create_node(params: dict) -> Node` function
|
||||
- A `node` variable of type Node
|
||||
- A class that inherits from Node
|
||||
|
||||
:param file_path: Path to the Python file.
|
||||
:type file_path: Path
|
||||
:returns: LoadedNode with the loaded node or error info.
|
||||
:rtype: LoadedNode
|
||||
"""
|
||||
file_hash = self._compute_file_hash(file_path)
|
||||
module_name = f"_dynamic_node_{file_path.stem}_{file_hash[:8]}"
|
||||
|
||||
try:
|
||||
# Load the module dynamically
|
||||
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise ImportError(f"Cannot load spec for {file_path}")
|
||||
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
# Try different ways to get the node
|
||||
node: Node | None = None
|
||||
|
||||
# Option 1: create_node function
|
||||
if hasattr(module, "create_node"):
|
||||
node = module.create_node(self.node_params)
|
||||
if not isinstance(node, Node):
|
||||
raise TypeError(
|
||||
f"create_node() must return a Node, got {type(node)}"
|
||||
)
|
||||
|
||||
# Option 2: node variable
|
||||
elif hasattr(module, "node"):
|
||||
node = module.node
|
||||
if not isinstance(node, Node):
|
||||
raise TypeError(f"'node' must be a Node instance, got {type(node)}")
|
||||
|
||||
# Option 3: Node subclass
|
||||
else:
|
||||
node_classes = [
|
||||
obj
|
||||
for name, obj in vars(module).items()
|
||||
if isinstance(obj, type)
|
||||
and issubclass(obj, Node)
|
||||
and obj is not Node
|
||||
and not name.startswith("_")
|
||||
]
|
||||
if node_classes:
|
||||
# Use the first non-private Node subclass found
|
||||
node = node_classes[0](
|
||||
f=lambda params, **kwargs: kwargs, # Default passthrough
|
||||
requires=[],
|
||||
provides=[],
|
||||
params=self.node_params,
|
||||
name=file_path.stem,
|
||||
)
|
||||
|
||||
if node is None:
|
||||
raise ValueError(
|
||||
f"No node found in {file_path}. "
|
||||
"Define create_node(params), a 'node' variable, or a Node subclass."
|
||||
)
|
||||
|
||||
# Register HTTP routes for HttpNode trigger nodes
|
||||
if isinstance(node, HttpNode) and node.mode == HttpNode.Mode.TRIGGER:
|
||||
if self._fastapi_app is not None:
|
||||
node.register_route(self._fastapi_app)
|
||||
else:
|
||||
logger.warning(
|
||||
f"HttpNode '{node.name}' is a trigger but no FastAPI app provided"
|
||||
)
|
||||
|
||||
# Track MQTT subscriber nodes for later subscription start
|
||||
# (subscriptions are started in _rebuild_pipeline after pipeline is built)
|
||||
if isinstance(node, MqttNode) and node.mode == MqttNode.Mode.SUBSCRIBER:
|
||||
logger.info(
|
||||
f"MqttNode '{node.name}' is a subscriber, will start after pipeline rebuild"
|
||||
)
|
||||
|
||||
# Track DelayNode cron jobs for later start
|
||||
# (cron jobs are started in _rebuild_pipeline after pipeline is built)
|
||||
if isinstance(node, DelayNode) and node.cron_expr:
|
||||
logger.info(
|
||||
f"DelayNode '{node.name}' has cron '{node.cron_expr}', will start after pipeline rebuild"
|
||||
)
|
||||
|
||||
logger.info(f"Loaded node '{node.name}' from {file_path.name}")
|
||||
|
||||
return LoadedNode(
|
||||
node=node,
|
||||
file_path=file_path,
|
||||
file_hash=file_hash,
|
||||
status=NodeStatus.ACTIVE,
|
||||
last_modified=file_path.stat().st_mtime,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
|
||||
logger.error(f"Failed to load node from {file_path.name}: {e}")
|
||||
|
||||
return LoadedNode(
|
||||
node=None,
|
||||
file_path=file_path,
|
||||
file_hash=file_hash,
|
||||
status=NodeStatus.ERROR,
|
||||
error_message=error_msg,
|
||||
last_modified=file_path.stat().st_mtime if file_path.exists() else 0,
|
||||
)
|
||||
finally:
|
||||
# Clean up the module from sys.modules to allow reloading
|
||||
if module_name in sys.modules:
|
||||
del sys.modules[module_name]
|
||||
|
||||
def _unload_node(self, node_id: str) -> None:
|
||||
"""
|
||||
Remove a node from the controller state (sync version).
|
||||
|
||||
:param node_id: The ID of the node to unload.
|
||||
:type node_id: str
|
||||
"""
|
||||
if node_id in self._state.nodes:
|
||||
loaded_node = self._state.nodes.pop(node_id)
|
||||
|
||||
# Unregister HTTP routes for HttpNode trigger nodes
|
||||
if (
|
||||
loaded_node.node is not None
|
||||
and isinstance(loaded_node.node, HttpNode)
|
||||
and loaded_node.node.mode == HttpNode.Mode.TRIGGER
|
||||
and self._fastapi_app is not None
|
||||
):
|
||||
loaded_node.node.unregister_route(self._fastapi_app)
|
||||
|
||||
# Note: MQTT subscriptions need async cleanup, handled in _unload_node_async
|
||||
|
||||
logger.info(f"Unloaded node '{node_id}'")
|
||||
for callback in self._on_node_removed:
|
||||
try:
|
||||
callback(node_id)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in on_node_removed callback: {e}")
|
||||
|
||||
async def _unload_node_async(self, node_id: str) -> None:
|
||||
"""
|
||||
Remove a node from the controller state (async version).
|
||||
|
||||
This version properly handles async cleanup like stopping MQTT subscriptions.
|
||||
|
||||
:param node_id: The ID of the node to unload.
|
||||
:type node_id: str
|
||||
"""
|
||||
if node_id in self._state.nodes:
|
||||
loaded_node = self._state.nodes.pop(node_id)
|
||||
|
||||
# Unregister HTTP routes for HttpNode trigger nodes
|
||||
if (
|
||||
loaded_node.node is not None
|
||||
and isinstance(loaded_node.node, HttpNode)
|
||||
and loaded_node.node.mode == HttpNode.Mode.TRIGGER
|
||||
and self._fastapi_app is not None
|
||||
):
|
||||
loaded_node.node.unregister_route(self._fastapi_app)
|
||||
|
||||
# Stop MQTT subscriptions for MqttNode subscriber nodes
|
||||
if (
|
||||
loaded_node.node is not None
|
||||
and isinstance(loaded_node.node, MqttNode)
|
||||
and loaded_node.node.mode == MqttNode.Mode.SUBSCRIBER
|
||||
):
|
||||
try:
|
||||
await loaded_node.node.stop_subscription()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error stopping MQTT subscription for '{node_id}': {e}"
|
||||
)
|
||||
|
||||
# Stop cron jobs for DelayNode nodes
|
||||
if (
|
||||
loaded_node.node is not None
|
||||
and isinstance(loaded_node.node, DelayNode)
|
||||
and loaded_node.node.cron_expr
|
||||
):
|
||||
try:
|
||||
await loaded_node.node.stop_cron()
|
||||
except Exception as e:
|
||||
logger.error(f"Error stopping cron for '{node_id}': {e}")
|
||||
|
||||
logger.info(f"Unloaded node '{node_id}'")
|
||||
for callback in self._on_node_removed:
|
||||
try:
|
||||
callback(node_id)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in on_node_removed callback: {e}")
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Pipeline Building
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
async def _rebuild_pipeline(self) -> None:
|
||||
"""
|
||||
Rebuild the pipeline from currently loaded nodes.
|
||||
|
||||
Only includes nodes with ACTIVE status. The pipeline is rebuilt
|
||||
atomically to avoid inconsistent states. Also starts MQTT subscriptions
|
||||
for subscriber nodes.
|
||||
"""
|
||||
async with self._rebuild_lock:
|
||||
active_nodes = [
|
||||
ln.node
|
||||
for ln in self._state.nodes.values()
|
||||
if ln.status == NodeStatus.ACTIVE and ln.node is not None
|
||||
]
|
||||
|
||||
if not active_nodes:
|
||||
logger.warning("No active nodes to build pipeline")
|
||||
self._state.pipeline = None
|
||||
return
|
||||
|
||||
try:
|
||||
# Build new pipeline
|
||||
new_pipeline = Pipeline(
|
||||
nodes=active_nodes,
|
||||
max_workers=self.max_workers,
|
||||
state=self.state_backend,
|
||||
)
|
||||
|
||||
self._state.pipeline = new_pipeline
|
||||
logger.info(f"Pipeline rebuilt with {len(active_nodes)} nodes")
|
||||
|
||||
# Start MQTT subscriptions for subscriber nodes
|
||||
for node in active_nodes:
|
||||
if (
|
||||
isinstance(node, MqttNode)
|
||||
and node.mode == MqttNode.Mode.SUBSCRIBER
|
||||
):
|
||||
if not node.is_subscribed:
|
||||
try:
|
||||
await node.start_subscription()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to start MQTT subscription for '{node.name}': {e}"
|
||||
)
|
||||
|
||||
# Start cron jobs for DelayNode nodes
|
||||
for node in active_nodes:
|
||||
if isinstance(node, DelayNode) and node.cron_expr:
|
||||
if node._cron_task is None:
|
||||
try:
|
||||
await node.start_cron()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to start cron for '{node.name}': {e}"
|
||||
)
|
||||
|
||||
# Notify callbacks
|
||||
for callback in self._on_pipeline_rebuilt:
|
||||
try:
|
||||
callback(new_pipeline)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in on_pipeline_rebuilt callback: {e}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to rebuild pipeline: {e}")
|
||||
# Keep the old pipeline intact
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# File Watching
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
async def _scan_nodes_directory(self) -> None:
|
||||
"""Scan the nodes directory and load all Python files."""
|
||||
if not self.nodes_dir.exists():
|
||||
logger.warning(f"Nodes directory does not exist: {self.nodes_dir}")
|
||||
self.nodes_dir.mkdir(parents=True, exist_ok=True)
|
||||
return
|
||||
|
||||
python_files = list(self.nodes_dir.glob("*.py"))
|
||||
logger.info(f"Found {len(python_files)} Python files in {self.nodes_dir}")
|
||||
|
||||
for file_path in python_files:
|
||||
if file_path.name.startswith("_"):
|
||||
continue # Skip private/internal files
|
||||
|
||||
node_id = file_path.stem
|
||||
loaded_node = self._load_node_from_file(file_path)
|
||||
self._state.nodes[node_id] = loaded_node
|
||||
|
||||
# Trigger callbacks
|
||||
if loaded_node.status == NodeStatus.ACTIVE and loaded_node.node:
|
||||
for callback in self._on_node_loaded:
|
||||
try:
|
||||
callback(node_id, loaded_node.node)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in on_node_loaded callback: {e}")
|
||||
elif loaded_node.status == NodeStatus.ERROR:
|
||||
for callback in self._on_node_error:
|
||||
try:
|
||||
callback(node_id, loaded_node.error_message or "Unknown error")
|
||||
except Exception as e:
|
||||
logger.error(f"Error in on_node_error callback: {e}")
|
||||
|
||||
await self._rebuild_pipeline()
|
||||
|
||||
async def _handle_file_change(self, change_type: Change, file_path: Path) -> None:
|
||||
"""Handle a file change event."""
|
||||
if not file_path.suffix == ".py" or file_path.name.startswith("_"):
|
||||
return
|
||||
|
||||
node_id = file_path.stem
|
||||
|
||||
if change_type == Change.deleted:
|
||||
await self._unload_node_async(node_id)
|
||||
await self._rebuild_pipeline()
|
||||
return
|
||||
|
||||
# Added or modified
|
||||
if node_id in self._state.nodes:
|
||||
# Check if file actually changed
|
||||
current_hash = self._state.nodes[node_id].file_hash
|
||||
try:
|
||||
new_hash = self._compute_file_hash(file_path)
|
||||
if new_hash == current_hash:
|
||||
return # No actual change
|
||||
except FileNotFoundError:
|
||||
return
|
||||
|
||||
# Unload existing node before reloading (for proper cleanup)
|
||||
await self._unload_node_async(node_id)
|
||||
|
||||
# Load/reload the node
|
||||
loaded_node = self._load_node_from_file(file_path)
|
||||
self._state.nodes[node_id] = loaded_node
|
||||
|
||||
# Trigger callbacks
|
||||
if loaded_node.status == NodeStatus.ACTIVE and loaded_node.node:
|
||||
for callback in self._on_node_loaded:
|
||||
try:
|
||||
callback(node_id, loaded_node.node)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in on_node_loaded callback: {e}")
|
||||
elif loaded_node.status == NodeStatus.ERROR:
|
||||
for callback in self._on_node_error:
|
||||
try:
|
||||
callback(node_id, loaded_node.error_message or "Unknown error")
|
||||
except Exception as e:
|
||||
logger.error(f"Error in on_node_error callback: {e}")
|
||||
|
||||
await self._rebuild_pipeline()
|
||||
|
||||
async def _watch_files(self) -> None:
|
||||
"""Watch the nodes directory for changes."""
|
||||
logger.info(f"Starting file watcher for {self.nodes_dir}")
|
||||
|
||||
try:
|
||||
async for changes in awatch(
|
||||
self.nodes_dir, stop_event=self._shutdown_event
|
||||
):
|
||||
for change_type, path_str in changes:
|
||||
file_path = Path(path_str)
|
||||
logger.info(f"File change detected: {change_type.name} {file_path}")
|
||||
await self._handle_file_change(change_type, file_path)
|
||||
except asyncio.CancelledError:
|
||||
logger.info("File watcher cancelled")
|
||||
except Exception as e:
|
||||
logger.error(f"File watcher error: {e}")
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Public API
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
async def start(self) -> None:
|
||||
"""
|
||||
Start the pipeline controller.
|
||||
|
||||
Performs initial scan of the nodes directory and starts the file watcher.
|
||||
"""
|
||||
if self._state.is_running:
|
||||
logger.warning("Controller is already running")
|
||||
return
|
||||
|
||||
self._state.is_running = True
|
||||
self._shutdown_event.clear()
|
||||
|
||||
# Initial scan
|
||||
await self._scan_nodes_directory()
|
||||
|
||||
# Start file watcher
|
||||
self._watch_task = asyncio.create_task(self._watch_files())
|
||||
logger.info("Pipeline controller started")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the pipeline controller, file watcher, and all subscriptions."""
|
||||
if not self._state.is_running:
|
||||
return
|
||||
|
||||
logger.info("Stopping pipeline controller...")
|
||||
self._shutdown_event.set()
|
||||
|
||||
# Stop all MQTT subscriptions
|
||||
for node_id, loaded_node in list(self._state.nodes.items()):
|
||||
if (
|
||||
loaded_node.node is not None
|
||||
and isinstance(loaded_node.node, MqttNode)
|
||||
and loaded_node.node.mode == MqttNode.Mode.SUBSCRIBER
|
||||
):
|
||||
try:
|
||||
await loaded_node.node.stop_subscription()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error stopping MQTT subscription for '{node_id}': {e}"
|
||||
)
|
||||
|
||||
# Stop all cron jobs
|
||||
for node_id, loaded_node in list(self._state.nodes.items()):
|
||||
if (
|
||||
loaded_node.node is not None
|
||||
and isinstance(loaded_node.node, DelayNode)
|
||||
and loaded_node.node.cron_expr
|
||||
):
|
||||
try:
|
||||
await loaded_node.node.stop_cron()
|
||||
except Exception as e:
|
||||
logger.error(f"Error stopping cron for '{node_id}': {e}")
|
||||
|
||||
if self._watch_task:
|
||||
self._watch_task.cancel()
|
||||
try:
|
||||
await self._watch_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._watch_task = None
|
||||
|
||||
self._state.is_running = False
|
||||
logger.info("Pipeline controller stopped")
|
||||
|
||||
async def reload_node(self, node_id: str) -> bool:
|
||||
"""
|
||||
Manually reload a specific node.
|
||||
|
||||
:param node_id: The ID (filename stem) of the node to reload.
|
||||
:type node_id: str
|
||||
:returns: True if reload was successful.
|
||||
:rtype: bool
|
||||
"""
|
||||
file_path = self.nodes_dir / f"{node_id}.py"
|
||||
if not file_path.exists():
|
||||
logger.error(f"Node file not found: {file_path}")
|
||||
return False
|
||||
|
||||
await self._handle_file_change(Change.modified, file_path)
|
||||
return (
|
||||
self._state.nodes.get(
|
||||
node_id,
|
||||
LoadedNode(
|
||||
node=None,
|
||||
file_path=file_path,
|
||||
file_hash="",
|
||||
status=NodeStatus.ERROR,
|
||||
),
|
||||
).status
|
||||
== NodeStatus.ACTIVE
|
||||
)
|
||||
|
||||
async def reload_all(self) -> None:
|
||||
"""Reload all nodes from the nodes directory."""
|
||||
self._state.nodes.clear()
|
||||
await self._scan_nodes_directory()
|
||||
|
||||
def add_node_from_code(self, node_id: str, code: str) -> bool:
|
||||
"""
|
||||
Add or update a node from code string.
|
||||
|
||||
This is the entry point for frontend-provided node code.
|
||||
|
||||
:param node_id: Unique identifier for the node.
|
||||
:type node_id: str
|
||||
:param code: Python code defining the node.
|
||||
:type code: str
|
||||
:returns: True if node was successfully added.
|
||||
:rtype: bool
|
||||
"""
|
||||
file_path = self.nodes_dir / f"{node_id}.py"
|
||||
|
||||
try:
|
||||
# Write the code to file
|
||||
file_path.write_text(code)
|
||||
logger.info(f"Wrote node code to {file_path}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to write node file: {e}")
|
||||
return False
|
||||
|
||||
def remove_node(self, node_id: str) -> bool:
|
||||
"""
|
||||
Remove a node by deleting its file.
|
||||
|
||||
:param node_id: The ID of the node to remove.
|
||||
:type node_id: str
|
||||
:returns: True if node was removed.
|
||||
:rtype: bool
|
||||
"""
|
||||
file_path = self.nodes_dir / f"{node_id}.py"
|
||||
|
||||
try:
|
||||
if file_path.exists():
|
||||
file_path.unlink()
|
||||
logger.info(f"Removed node file {file_path}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to remove node file: {e}")
|
||||
return False
|
||||
|
||||
@property
|
||||
def pipeline(self) -> Pipeline | None:
|
||||
"""Get the current pipeline instance."""
|
||||
return self._state.pipeline
|
||||
|
||||
@property
|
||||
def nodes(self) -> dict[str, LoadedNode]:
|
||||
"""Get all loaded nodes with their status."""
|
||||
return self._state.nodes.copy()
|
||||
|
||||
@property
|
||||
def active_nodes(self) -> list[Node]:
|
||||
"""Get list of currently active nodes."""
|
||||
return [
|
||||
ln.node
|
||||
for ln in self._state.nodes.values()
|
||||
if ln.status == NodeStatus.ACTIVE and ln.node is not None
|
||||
]
|
||||
|
||||
@property
|
||||
def error_nodes(self) -> dict[str, str]:
|
||||
"""Get nodes that failed to load with their error messages."""
|
||||
return {
|
||||
node_id: ln.error_message or "Unknown error"
|
||||
for node_id, ln in self._state.nodes.items()
|
||||
if ln.status == NodeStatus.ERROR
|
||||
}
|
||||
|
||||
def get_node_status(self, node_id: str) -> NodeStatus | None:
|
||||
"""Get the status of a specific node."""
|
||||
if node_id in self._state.nodes:
|
||||
return self._state.nodes[node_id].status
|
||||
return None
|
||||
|
||||
def trigger_node(self, node_id: str, inputs: dict | None = None) -> Any:
|
||||
"""
|
||||
Trigger a node by its ID.
|
||||
|
||||
:param node_id: The ID of the node to trigger.
|
||||
:type node_id: str
|
||||
:param inputs: Optional input values.
|
||||
:type inputs: dict | None
|
||||
:returns: Result of the node execution.
|
||||
:rtype: Any
|
||||
"""
|
||||
if self._state.pipeline is None:
|
||||
raise RuntimeError("Pipeline not initialized")
|
||||
|
||||
loaded_node = self._state.nodes.get(node_id)
|
||||
if loaded_node is None or loaded_node.node is None:
|
||||
raise ValueError(f"Node '{node_id}' not found or not active")
|
||||
|
||||
return loaded_node.node.trigger(inputs)
|
||||
|
||||
def set_fastapi_app(self, app: "FastAPI") -> None:
|
||||
"""
|
||||
Set the FastAPI application for HTTP node route registration.
|
||||
|
||||
This method can be called after initialization to attach a FastAPI
|
||||
app. It will register routes for any existing HTTP trigger nodes.
|
||||
|
||||
:param app: The FastAPI application instance.
|
||||
:type app: FastAPI
|
||||
"""
|
||||
self._fastapi_app = app
|
||||
|
||||
# Register routes for any existing HTTP trigger nodes
|
||||
for node_id, loaded_node in self._state.nodes.items():
|
||||
if (
|
||||
loaded_node.node is not None
|
||||
and loaded_node.status == NodeStatus.ACTIVE
|
||||
and isinstance(loaded_node.node, HttpNode)
|
||||
and loaded_node.node.mode == HttpNode.Mode.TRIGGER
|
||||
):
|
||||
loaded_node.node.register_route(app)
|
||||
logger.info(f"Registered HTTP route for existing node '{node_id}'")
|
||||
|
||||
|
||||
async def main():
|
||||
"""
|
||||
Main function to demonstrate the PipelineController usage.
|
||||
|
||||
Sets up Redis or falls back to MemoryState, creates a controller,
|
||||
and watches for node changes in the nodes directory.
|
||||
"""
|
||||
try:
|
||||
state = RedisState(
|
||||
host="localhost",
|
||||
port=6379,
|
||||
namespace="fluksio_controller_test",
|
||||
ttl=3600,
|
||||
)
|
||||
if state.ping():
|
||||
logger.info("Using Redis state backend")
|
||||
else:
|
||||
raise ConnectionError("Redis ping failed")
|
||||
except Exception as e:
|
||||
logger.warning("Redis not available (%s), using MemoryState", e)
|
||||
state = MemoryState()
|
||||
|
||||
# Create controller pointing to example nodes
|
||||
nodes_dir = Path(__file__).parent / "nodes"
|
||||
|
||||
controller = PipelineController(
|
||||
nodes_dir=nodes_dir,
|
||||
state_backend=state,
|
||||
max_workers=4,
|
||||
)
|
||||
|
||||
# Register callbacks to see what's happening
|
||||
controller.on_node_loaded(
|
||||
lambda nid, node: logger.info(
|
||||
"Loaded node '%s': %d inputs, %d outputs",
|
||||
nid,
|
||||
len(node.requires),
|
||||
len(node.provides),
|
||||
)
|
||||
)
|
||||
controller.on_node_error(
|
||||
lambda nid, err: logger.error("Error in node '%s': %s", nid, err.split("\n")[0])
|
||||
)
|
||||
controller.on_node_removed(lambda nid: logger.info("Removed node '%s'", nid))
|
||||
controller.on_pipeline_rebuilt(
|
||||
lambda p: logger.info(
|
||||
"Pipeline rebuilt: %d nodes, %d dependencies",
|
||||
len(p.nodes),
|
||||
len(p.dependencies),
|
||||
)
|
||||
)
|
||||
|
||||
await controller.start()
|
||||
|
||||
# Show current state
|
||||
logger.info("--- Current State ---")
|
||||
logger.info("Active nodes: %s", [n.name for n in controller.active_nodes])
|
||||
logger.info("Error nodes: %s", list(controller.error_nodes.keys()))
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(1)
|
||||
controller.pipeline.visualize()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
asyncio.run(main())
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Main process interrupted by user")
|
||||
@@ -0,0 +1,218 @@
|
||||
"""
|
||||
DAG generation utilities for testing pipelines.
|
||||
|
||||
This module provides functions to generate random directed acyclic graphs (DAGs)
|
||||
with guaranteed properties like connectivity and no cycles.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import networkx as nx
|
||||
import numpy as np
|
||||
from typing import Callable
|
||||
|
||||
from util import Message, NodeParams
|
||||
from nodes import Node, MLPNode
|
||||
|
||||
|
||||
def generate_dag(
|
||||
num_nodes: int,
|
||||
edge_probability: float = 0.3,
|
||||
seed: int | None = None,
|
||||
) -> nx.DiGraph:
|
||||
"""
|
||||
Generate a random DAG using the Erdős-Rényi model with topological ordering.
|
||||
|
||||
Creates edges only from lower-indexed to higher-indexed nodes to guarantee
|
||||
acyclicity, then removes isolated nodes to ensure connectivity.
|
||||
|
||||
:param num_nodes: Number of nodes in the graph.
|
||||
:type num_nodes: int
|
||||
:param edge_probability: Probability of edge between any two nodes.
|
||||
:type edge_probability: float
|
||||
:param seed: Random seed for reproducibility.
|
||||
:type seed: int | None
|
||||
:returns: A random DAG.
|
||||
:rtype: nx.DiGraph
|
||||
"""
|
||||
rng = np.random.default_rng(seed)
|
||||
G = nx.DiGraph()
|
||||
G.add_nodes_from(range(num_nodes))
|
||||
|
||||
# Add edges only from lower to higher index (guarantees DAG)
|
||||
for i in range(num_nodes):
|
||||
for j in range(i + 1, num_nodes):
|
||||
if rng.random() < edge_probability:
|
||||
G.add_edge(i, j)
|
||||
|
||||
# Ensure connectivity: connect isolated nodes
|
||||
_ensure_connected(G, rng)
|
||||
|
||||
return G
|
||||
|
||||
|
||||
def _ensure_connected(G: nx.DiGraph, rng: np.random.Generator) -> None:
|
||||
"""
|
||||
Ensure the DAG is weakly connected by adding edges.
|
||||
|
||||
:param G: The graph to modify in place.
|
||||
:type G: nx.DiGraph
|
||||
:param rng: Random number generator.
|
||||
:type rng: np.random.Generator
|
||||
"""
|
||||
# Get weakly connected components
|
||||
components = list(nx.weakly_connected_components(G))
|
||||
|
||||
if len(components) <= 1:
|
||||
return
|
||||
|
||||
# Sort nodes in each component by index for DAG-safe edge addition
|
||||
sorted_components = [sorted(c) for c in components]
|
||||
|
||||
# Connect components by adding edge from max of one to min of next
|
||||
for i in range(len(sorted_components) - 1):
|
||||
src = sorted_components[i][-1] # Last (highest) node in component
|
||||
dst = sorted_components[i + 1][0] # First (lowest) node in next component
|
||||
|
||||
# Ensure edge direction maintains DAG property
|
||||
if src < dst:
|
||||
G.add_edge(src, dst)
|
||||
else:
|
||||
G.add_edge(dst, src)
|
||||
|
||||
|
||||
def dag_to_pipeline_nodes(
|
||||
G: nx.DiGraph,
|
||||
node_factory: Callable[..., Node] | None = None,
|
||||
params: dict | None = None,
|
||||
seed: int | None = None,
|
||||
) -> tuple[list[Node], list[Message], list[Node]]:
|
||||
"""
|
||||
Convert a networkx DAG to pipeline nodes with messages.
|
||||
|
||||
:param G: The DAG to convert.
|
||||
:type G: nx.DiGraph
|
||||
:param node_factory: Factory function to create nodes (default: MLPNode).
|
||||
:type node_factory: Callable[..., Node] | None
|
||||
:param params: Parameters to pass to node factory.
|
||||
:type params: dict | None
|
||||
:param seed: Random seed for reproducibility.
|
||||
:type seed: int | None
|
||||
:returns: Tuple of (nodes, all_messages, trigger_nodes).
|
||||
:rtype: tuple[list[Node], list[Message], list[Node]]
|
||||
"""
|
||||
rng = np.random.default_rng(seed)
|
||||
params = params or {}
|
||||
|
||||
# Create a message for each edge
|
||||
edge_messages: dict[tuple[int, int], Message] = {}
|
||||
for src, dst in G.edges():
|
||||
edge_messages[(src, dst)] = Message(name=f"msg_{src}_{dst}")
|
||||
|
||||
# Create nodes in topological order
|
||||
nodes: list[Node] = []
|
||||
trigger_nodes: list[Node] = []
|
||||
|
||||
for node_id in nx.topological_sort(G):
|
||||
# Inputs: messages from incoming edges
|
||||
requires = [edge_messages[(src, node_id)] for src in G.predecessors(node_id)]
|
||||
|
||||
# Outputs: messages for outgoing edges
|
||||
provides = [edge_messages[(node_id, dst)] for dst in G.successors(node_id)]
|
||||
|
||||
# Skip nodes with no outputs (sink nodes produce no messages)
|
||||
# But we still need to create them to consume inputs
|
||||
if node_factory:
|
||||
node = node_factory(
|
||||
requires=requires,
|
||||
provides=provides,
|
||||
params=params,
|
||||
name=f"{node_id}",
|
||||
)
|
||||
else:
|
||||
node = MLPNode(
|
||||
requires=requires,
|
||||
provides=provides,
|
||||
params=params,
|
||||
name=f"{node_id}",
|
||||
)
|
||||
|
||||
nodes.append(node)
|
||||
|
||||
# Trigger nodes are sources (no inputs)
|
||||
if len(requires) == 0:
|
||||
trigger_nodes.append(node)
|
||||
|
||||
all_messages = list(edge_messages.values())
|
||||
return nodes, all_messages, trigger_nodes
|
||||
|
||||
|
||||
def generate_multi_pipeline_dag(
|
||||
num_pipelines: int,
|
||||
nodes_per_pipeline: int,
|
||||
cross_pipeline_edges: int = 2,
|
||||
edge_probability: float = 0.3,
|
||||
seed: int | None = None,
|
||||
) -> tuple[list[nx.DiGraph], nx.DiGraph]:
|
||||
"""
|
||||
Generate multiple DAGs that can be composed into a parent pipeline.
|
||||
|
||||
Creates separate DAGs for each pipeline and adds cross-pipeline edges
|
||||
to create dependencies between them.
|
||||
|
||||
:param num_pipelines: Number of child pipelines.
|
||||
:type num_pipelines: int
|
||||
:param nodes_per_pipeline: Nodes in each pipeline.
|
||||
:type nodes_per_pipeline: int
|
||||
:param cross_pipeline_edges: Number of edges connecting pipelines.
|
||||
:type cross_pipeline_edges: int
|
||||
:param edge_probability: Edge probability within each pipeline.
|
||||
:type edge_probability: float
|
||||
:param seed: Random seed for reproducibility.
|
||||
:type seed: int | None
|
||||
:returns: Tuple of (list of pipeline DAGs, combined DAG).
|
||||
:rtype: tuple[list[nx.DiGraph], nx.DiGraph]
|
||||
"""
|
||||
rng = np.random.default_rng(seed)
|
||||
|
||||
pipeline_graphs: list[nx.DiGraph] = []
|
||||
combined = nx.DiGraph()
|
||||
|
||||
# Generate each pipeline's DAG
|
||||
for p in range(num_pipelines):
|
||||
offset = p * nodes_per_pipeline
|
||||
G = generate_dag(
|
||||
nodes_per_pipeline, edge_probability, seed=rng.integers(0, 10000)
|
||||
)
|
||||
|
||||
# Relabel nodes with offset
|
||||
mapping = {n: n + offset for n in G.nodes()}
|
||||
G = nx.relabel_nodes(G, mapping)
|
||||
|
||||
# Add pipeline attribute
|
||||
for n in G.nodes():
|
||||
G.nodes[n]["pipeline"] = p
|
||||
|
||||
pipeline_graphs.append(G)
|
||||
combined = nx.compose(combined, G)
|
||||
|
||||
# Add cross-pipeline edges (from later nodes to earlier nodes of next pipeline)
|
||||
for _ in range(cross_pipeline_edges):
|
||||
p1, p2 = rng.choice(num_pipelines, size=2, replace=False)
|
||||
if p1 > p2:
|
||||
p1, p2 = p2, p1
|
||||
|
||||
# Get sink nodes from p1 (nodes with no outgoing edges within pipeline)
|
||||
p1_nodes = [n for n in pipeline_graphs[p1].nodes()]
|
||||
p1_sinks = [n for n in p1_nodes if pipeline_graphs[p1].out_degree(n) == 0]
|
||||
|
||||
# Get source-ish nodes from p2 (nodes with few incoming edges)
|
||||
p2_nodes = [n for n in pipeline_graphs[p2].nodes()]
|
||||
p2_sources = [n for n in p2_nodes if combined.in_degree(n) <= 1]
|
||||
|
||||
if p1_sinks and p2_sources:
|
||||
src = rng.choice(p1_sinks)
|
||||
dst = rng.choice(p2_sources)
|
||||
combined.add_edge(src, dst)
|
||||
|
||||
return pipeline_graphs, combined
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,20 @@
|
||||
from nodes import DelayNode
|
||||
from util import Message
|
||||
import numpy as np
|
||||
import logging
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_node(params: dict) -> DelayNode:
|
||||
|
||||
return DelayNode(
|
||||
requires=[],
|
||||
provides=[Message(name="timestamp", dtype=float)],
|
||||
params={"cron": "* * * * * *"},
|
||||
name="alarm",
|
||||
)
|
||||
@@ -0,0 +1,25 @@
|
||||
from nodes import DelayNode
|
||||
from util import Message
|
||||
|
||||
import logging
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_node(params: dict) -> DelayNode:
|
||||
return DelayNode(
|
||||
params={
|
||||
"delay": 1,
|
||||
},
|
||||
requires=[
|
||||
Message(name="random_value_rec_rate", dtype=int),
|
||||
],
|
||||
provides=[
|
||||
Message(name="random_value_rec_rate_delayed", dtype=int),
|
||||
],
|
||||
name="delay",
|
||||
)
|
||||
@@ -0,0 +1,35 @@
|
||||
from nodes import InfluxDbNode
|
||||
from util import Message
|
||||
|
||||
import logging
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_node(params: dict) -> InfluxDbNode:
|
||||
return InfluxDbNode(
|
||||
requires=[
|
||||
Message(name="random_value_rec_rate", dtype=int),
|
||||
],
|
||||
params={
|
||||
"url": params.get("influxdb_url", "http://10.200.200.115:8086"),
|
||||
"token": params.get(
|
||||
"influxdb_token",
|
||||
"***REMOVED-INFLUXDB-TOKEN***==",
|
||||
),
|
||||
"org": params.get("influxdb_org", "strobl"),
|
||||
"bucket": params.get("influxdb_bucket", "test"),
|
||||
"writes": {
|
||||
"random_value_rec_rate": {
|
||||
"measurement": "test",
|
||||
"field": "random",
|
||||
},
|
||||
},
|
||||
"synchronous": True,
|
||||
},
|
||||
name="influxdb",
|
||||
)
|
||||
@@ -0,0 +1,25 @@
|
||||
from nodes import MqttNode
|
||||
from util import Message
|
||||
|
||||
import logging
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_node(params: dict) -> MqttNode:
|
||||
return MqttNode(
|
||||
requires=[
|
||||
Message(name="random_value_send", dtype=int),
|
||||
],
|
||||
params={
|
||||
"broker_host": "127.0.0.1",
|
||||
"topic": {
|
||||
"random_value_send": "random_value",
|
||||
},
|
||||
},
|
||||
name="mqtt_a",
|
||||
)
|
||||
@@ -0,0 +1,25 @@
|
||||
from nodes import MqttNode
|
||||
from util import Message
|
||||
|
||||
import logging
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_node(params: dict) -> MqttNode:
|
||||
return MqttNode(
|
||||
provides=[
|
||||
Message(name="random_value_rec", dtype=int),
|
||||
],
|
||||
params={
|
||||
"broker_host": "127.0.0.1",
|
||||
"topic": {
|
||||
"random_value_rec": "random_value",
|
||||
},
|
||||
},
|
||||
name="mqtt_b",
|
||||
)
|
||||
@@ -0,0 +1,30 @@
|
||||
from nodes import Node
|
||||
from util import Message
|
||||
|
||||
import logging
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_node(params: dict) -> Node:
|
||||
def check_power(params, random_value_rec_rate_delayed: int, **kwargs):
|
||||
threshold = params.get("alert_threshold", 1000.0)
|
||||
|
||||
if random_value_rec_rate_delayed > 50:
|
||||
logger.info(f"[alert_node] ⚠️ Input >50")
|
||||
else:
|
||||
logger.info(f"[alert_node] ✓ Input < 50")
|
||||
|
||||
return Node(
|
||||
f=check_power,
|
||||
requires=[
|
||||
Message(name="random_value_rec_rate_delayed", dtype=int),
|
||||
],
|
||||
provides=[], # Using object for optional string
|
||||
params=params,
|
||||
name="notify",
|
||||
)
|
||||
@@ -0,0 +1,32 @@
|
||||
from nodes import Node
|
||||
from util import Message
|
||||
import numpy as np
|
||||
import logging
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_node(params: dict) -> Node:
|
||||
def f(params, timestamp):
|
||||
seed = int(params.get("seed", 1000))
|
||||
rng = np.random.default_rng(seed=seed)
|
||||
|
||||
rand_value = int(100 * rng.uniform())
|
||||
|
||||
return {"random_value_send": rand_value}
|
||||
|
||||
return Node(
|
||||
provides=[
|
||||
Message(name="random_value_send", dtype=int),
|
||||
],
|
||||
requires=[Message(name="timestamp", dtype=float)],
|
||||
f=f,
|
||||
params={
|
||||
"seed": "1000",
|
||||
},
|
||||
name="random",
|
||||
)
|
||||
@@ -0,0 +1,26 @@
|
||||
from nodes import DelayNode
|
||||
from util import Message
|
||||
|
||||
import logging
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_node(params: dict) -> DelayNode:
|
||||
return DelayNode(
|
||||
params={
|
||||
"delay": 0,
|
||||
"interval": 3,
|
||||
},
|
||||
requires=[
|
||||
Message(name="random_value_rec", dtype=int),
|
||||
],
|
||||
provides=[
|
||||
Message(name="random_value_rec_rate", dtype=int),
|
||||
],
|
||||
name="rate",
|
||||
)
|
||||
@@ -0,0 +1,44 @@
|
||||
"""
|
||||
Example sink node: Alert Generator
|
||||
|
||||
This node consumes the comfort index and generates alerts
|
||||
when comfort drops below a threshold. It demonstrates a
|
||||
sink node (produces no outputs for other nodes).
|
||||
"""
|
||||
|
||||
from nodes import Node
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params: dict) -> Node:
|
||||
"""Create the alert generator node."""
|
||||
|
||||
def check_comfort_alert(params, comfort_index: float = 50.0, **kwargs):
|
||||
"""
|
||||
Check if comfort level requires an alert.
|
||||
|
||||
Generates an alert message if comfort is too low.
|
||||
"""
|
||||
threshold = params.get("alert_threshold", 60.0)
|
||||
|
||||
if comfort_index < threshold:
|
||||
alert_msg = f"LOW COMFORT ALERT: Index {comfort_index:.1f} below threshold {threshold}"
|
||||
print(f"[alert_node] ⚠️ {alert_msg}")
|
||||
# In a real implementation, this could:
|
||||
# - Send email/SMS
|
||||
# - Push to a message queue
|
||||
# - Trigger home automation
|
||||
return {"alert": alert_msg}
|
||||
else:
|
||||
print(f"[alert_node] ✓ Comfort level OK ({comfort_index:.1f})")
|
||||
return {"alert": None}
|
||||
|
||||
return Node(
|
||||
f=check_comfort_alert,
|
||||
requires=[Message(name="comfort_index", dtype=float)],
|
||||
provides=[
|
||||
Message(name="alert", dtype=object)
|
||||
], # Using object for optional string
|
||||
params=params,
|
||||
name="alert_node",
|
||||
)
|
||||
@@ -0,0 +1,46 @@
|
||||
"""
|
||||
Example processing node: Comfort Calculator
|
||||
|
||||
This node takes temperature and humidity as inputs and calculates
|
||||
a comfort index. It demonstrates how nodes can have dependencies
|
||||
that are automatically resolved by the pipeline.
|
||||
"""
|
||||
|
||||
from nodes import Node
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params: dict) -> Node:
|
||||
"""Create the comfort calculator node."""
|
||||
|
||||
def calculate_comfort(
|
||||
params, temperature: float = 20.0, humidity: float = 50.0, **kwargs
|
||||
):
|
||||
"""
|
||||
Calculate comfort index based on temperature and humidity.
|
||||
|
||||
Uses a simplified heat index formula.
|
||||
"""
|
||||
# Simplified comfort calculation
|
||||
# Ideal: 22°C, 45% humidity
|
||||
temp_diff = abs(temperature - 22.0)
|
||||
humidity_diff = abs(humidity - 45.0)
|
||||
|
||||
comfort = 100.0 - (temp_diff * 3) - (humidity_diff * 0.5)
|
||||
comfort = max(0, min(100, comfort))
|
||||
|
||||
print(
|
||||
f"[comfort_calculator] T={temperature:.1f}°C, H={humidity:.1f}% -> Comfort={comfort:.1f}"
|
||||
)
|
||||
return {"comfort_index": comfort}
|
||||
|
||||
return Node(
|
||||
f=calculate_comfort,
|
||||
requires=[
|
||||
Message(name="temperature", dtype=float),
|
||||
Message(name="humidity", dtype=float),
|
||||
],
|
||||
provides=[Message(name="comfort_index", dtype=float)],
|
||||
params=params,
|
||||
name="comfort_calculator",
|
||||
)
|
||||
@@ -0,0 +1,35 @@
|
||||
"""
|
||||
Example HTTP sender node: Data Publisher
|
||||
|
||||
This node sends processed data to an external API endpoint
|
||||
via HTTP POST requests.
|
||||
"""
|
||||
|
||||
from nodes import HttpNode
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params: dict) -> HttpNode:
|
||||
"""
|
||||
Create an HTTP sender node that publishes comfort data.
|
||||
|
||||
:param params: Parameters passed from the controller.
|
||||
:type params: dict
|
||||
:returns: Configured HttpNode instance in sender mode.
|
||||
:rtype: HttpNode
|
||||
"""
|
||||
# Get the target URL from params or use a default
|
||||
target_url = params.get("webhook_url", "https://httpbin.org/post")
|
||||
|
||||
return HttpNode(
|
||||
url=target_url,
|
||||
method="POST",
|
||||
requires=[
|
||||
Message(name="comfort_index", dtype=float),
|
||||
Message(name="alert", dtype=object),
|
||||
],
|
||||
params=params,
|
||||
name="data_publisher",
|
||||
timeout=10.0,
|
||||
headers={"X-Source": "fluksio-pipeline"},
|
||||
)
|
||||
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
Example HTTP trigger node: Webhook Receiver
|
||||
|
||||
This node acts as a webhook endpoint that receives temperature data
|
||||
via HTTP POST requests and injects it into the pipeline.
|
||||
"""
|
||||
|
||||
from nodes import HttpNode
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params: dict) -> HttpNode:
|
||||
"""
|
||||
Create an HTTP trigger node that receives temperature data.
|
||||
|
||||
:param params: Parameters passed from the controller.
|
||||
:type params: dict
|
||||
:returns: Configured HttpNode instance in trigger mode.
|
||||
:rtype: HttpNode
|
||||
"""
|
||||
return HttpNode(
|
||||
url="/api/sensors/temperature",
|
||||
method="POST",
|
||||
provides=[
|
||||
Message(name="temperature", dtype=float),
|
||||
Message(name="sensor_id", dtype=str),
|
||||
],
|
||||
params=params,
|
||||
name="temperature_webhook",
|
||||
)
|
||||
@@ -0,0 +1,29 @@
|
||||
"""
|
||||
Example trigger node: Humidity Sensor
|
||||
|
||||
Another trigger node that simulates a humidity sensor.
|
||||
"""
|
||||
|
||||
from nodes import Node
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params: dict) -> Node:
|
||||
"""Create the humidity sensor node."""
|
||||
|
||||
def read_humidity(params, **kwargs):
|
||||
"""Simulate reading humidity from a sensor."""
|
||||
import random
|
||||
|
||||
humidity = 50.0 + random.gauss(0, 10) # 50% ± 10%
|
||||
humidity = max(0, min(100, humidity)) # Clamp to [0, 100]
|
||||
print(f"[humidity_sensor] Read humidity: {humidity:.1f}%")
|
||||
return {"humidity": humidity}
|
||||
|
||||
return Node(
|
||||
f=read_humidity,
|
||||
requires=[],
|
||||
provides=[Message(name="humidity", dtype=float)],
|
||||
params=params,
|
||||
name="humidity_sensor",
|
||||
)
|
||||
@@ -0,0 +1,52 @@
|
||||
"""
|
||||
Example InfluxDB reader node: Average Temperature Provider
|
||||
|
||||
This node queries InfluxDB for average temperature and provides it
|
||||
to downstream nodes.
|
||||
"""
|
||||
|
||||
from nodes import InfluxDbNode
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params: dict) -> InfluxDbNode:
|
||||
"""
|
||||
Create an InfluxDB reader node that queries average temperature.
|
||||
|
||||
:param params: Parameters passed from the controller, should include:
|
||||
- ``influxdb_url``: InfluxDB server URL
|
||||
- ``influxdb_token``: Authentication token
|
||||
- ``influxdb_org``: Organization name
|
||||
- ``influxdb_bucket``: Bucket name
|
||||
:type params: dict
|
||||
:returns: Configured InfluxDbNode instance.
|
||||
:rtype: InfluxDbNode
|
||||
"""
|
||||
return InfluxDbNode(
|
||||
provides=[
|
||||
Message(name="avg_temperature", dtype=float),
|
||||
Message(name="max_temperature", dtype=float),
|
||||
],
|
||||
params={
|
||||
"url": params.get("influxdb_url", "http://localhost:8086"),
|
||||
"token": params.get("influxdb_token", "my-token"),
|
||||
"org": params.get("influxdb_org", "my-org"),
|
||||
"bucket": params.get("influxdb_bucket", "sensors"),
|
||||
"query_range": "-1h",
|
||||
"queries": {
|
||||
"avg_temperature": {
|
||||
"measurement": "temperature",
|
||||
"field": "value",
|
||||
"tags": {"location": "room1"},
|
||||
"aggregation": "mean",
|
||||
},
|
||||
"max_temperature": {
|
||||
"measurement": "temperature",
|
||||
"field": "value",
|
||||
"tags": {"location": "room1"},
|
||||
"aggregation": "max",
|
||||
},
|
||||
},
|
||||
},
|
||||
name="influxdb_temperature_reader",
|
||||
)
|
||||
@@ -0,0 +1,54 @@
|
||||
"""
|
||||
Example InfluxDB writer node: Temperature Logger
|
||||
|
||||
This node writes temperature readings to InfluxDB.
|
||||
The input is just a float value - the measurement, field, and tags
|
||||
are configured via params.
|
||||
"""
|
||||
|
||||
from nodes import InfluxDbNode
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params: dict) -> InfluxDbNode:
|
||||
"""
|
||||
Create an InfluxDB writer node that logs temperature data.
|
||||
|
||||
The node receives a simple float value and writes it to InfluxDB
|
||||
with the configured measurement, field, and tags.
|
||||
|
||||
:param params: Parameters passed from the controller, should include:
|
||||
- ``influxdb_url``: InfluxDB server URL
|
||||
- ``influxdb_token``: Authentication token
|
||||
- ``influxdb_org``: Organization name
|
||||
- ``influxdb_bucket``: Bucket name
|
||||
:type params: dict
|
||||
:returns: Configured InfluxDbNode instance.
|
||||
:rtype: InfluxDbNode
|
||||
"""
|
||||
return InfluxDbNode(
|
||||
requires=[
|
||||
Message(name="temperature", dtype=float),
|
||||
Message(name="humidity", dtype=float),
|
||||
],
|
||||
params={
|
||||
"url": params.get("influxdb_url", "http://localhost:8086"),
|
||||
"token": params.get("influxdb_token", "my-token"),
|
||||
"org": params.get("influxdb_org", "my-org"),
|
||||
"bucket": params.get("influxdb_bucket", "sensors"),
|
||||
"write_precision": "ms",
|
||||
"writes": {
|
||||
"temperature": {
|
||||
"measurement": "environment",
|
||||
"field": "temp_celsius",
|
||||
"tags": {"location": "room1", "sensor": "dht22"},
|
||||
},
|
||||
"humidity": {
|
||||
"measurement": "environment",
|
||||
"field": "humidity_percent",
|
||||
"tags": {"location": "room1", "sensor": "dht22"},
|
||||
},
|
||||
},
|
||||
},
|
||||
name="influxdb_temperature_writer",
|
||||
)
|
||||
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
Example MQTT publisher node: Alert Publisher
|
||||
|
||||
This node publishes alerts to an MQTT topic when comfort levels
|
||||
are outside acceptable ranges.
|
||||
"""
|
||||
|
||||
from nodes import MqttNode
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params: dict) -> MqttNode:
|
||||
"""
|
||||
Create an MQTT publisher node that sends alert messages.
|
||||
|
||||
:param params: Parameters passed from the controller, should include:
|
||||
- ``broker_host``: MQTT broker hostname
|
||||
- ``broker_port``: MQTT broker port (optional, default 1883)
|
||||
- ``qos``: Quality of Service level (optional, default 0)
|
||||
:type params: dict
|
||||
:returns: Configured MqttNode instance in publisher mode.
|
||||
:rtype: MqttNode
|
||||
"""
|
||||
return MqttNode(
|
||||
topic="alerts/comfort",
|
||||
requires=[
|
||||
Message(name="comfort_index", dtype=float),
|
||||
Message(name="alert", dtype=object),
|
||||
],
|
||||
params={
|
||||
**params,
|
||||
"qos": params.get("qos", 1), # Use QoS 1 for alerts
|
||||
"retain": True, # Retain last alert
|
||||
},
|
||||
name="mqtt_alert_publisher",
|
||||
)
|
||||
@@ -0,0 +1,31 @@
|
||||
"""
|
||||
Example MQTT subscriber node: Temperature Sensor Listener
|
||||
|
||||
This node subscribes to an MQTT topic and triggers the pipeline
|
||||
when temperature readings are received.
|
||||
"""
|
||||
|
||||
from nodes import MqttNode
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params: dict) -> MqttNode:
|
||||
"""
|
||||
Create an MQTT subscriber node that listens for temperature data.
|
||||
|
||||
:param params: Parameters passed from the controller, should include:
|
||||
- ``broker_host``: MQTT broker hostname
|
||||
- ``broker_port``: MQTT broker port (optional, default 1883)
|
||||
:type params: dict
|
||||
:returns: Configured MqttNode instance in subscriber mode.
|
||||
:rtype: MqttNode
|
||||
"""
|
||||
return MqttNode(
|
||||
topic="sensors/temperature",
|
||||
provides=[
|
||||
Message(name="temperature", dtype=float),
|
||||
Message(name="sensor_id", dtype=str),
|
||||
],
|
||||
params=params,
|
||||
name="mqtt_temperature_listener",
|
||||
)
|
||||
@@ -0,0 +1,37 @@
|
||||
"""
|
||||
Example trigger node: Temperature Sensor
|
||||
|
||||
This is a trigger node (no inputs) that simulates a temperature sensor.
|
||||
Trigger nodes act as entry points to the pipeline - they can be triggered
|
||||
externally via HTTP, MQTT, or other mechanisms.
|
||||
"""
|
||||
|
||||
from nodes import Node
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params: dict) -> Node:
|
||||
"""
|
||||
Create the temperature sensor node.
|
||||
|
||||
:param params: Parameters passed from the controller.
|
||||
:type params: dict
|
||||
:returns: Configured Node instance.
|
||||
:rtype: Node
|
||||
"""
|
||||
|
||||
def read_temperature(params, **kwargs):
|
||||
"""Simulate reading temperature from a sensor."""
|
||||
import random
|
||||
|
||||
temperature = 20.0 + random.gauss(0, 2) # 20°C ± 2°C
|
||||
print(f"[temperature_sensor] Read temperature: {temperature:.2f}°C")
|
||||
return {"temperature": temperature}
|
||||
|
||||
return Node(
|
||||
f=read_temperature,
|
||||
requires=[], # No inputs - this is a trigger node
|
||||
provides=[Message(name="temperature", dtype=float)],
|
||||
params=params,
|
||||
name="temperature_sensor",
|
||||
)
|
||||
@@ -0,0 +1,626 @@
|
||||
"""
|
||||
Pipeline module for directed acyclic graph execution.
|
||||
|
||||
This module provides a Pipeline class that manages nodes with automatic
|
||||
dependency resolution and supports both sequential and parallel execution.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from concurrent.futures import ThreadPoolExecutor, Future, wait
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import networkx as nx
|
||||
|
||||
from nodes import Node
|
||||
from state import StateBackend, MemoryState
|
||||
from util import Message
|
||||
|
||||
|
||||
class Pipeline:
|
||||
"""
|
||||
Directed acyclic graph of nodes with automatic dependency resolution.
|
||||
|
||||
A Pipeline manages a collection of nodes, automatically resolving their
|
||||
dependencies and executing them in the correct order. Supports parallel
|
||||
execution, external triggering, and composition of multiple pipelines.
|
||||
|
||||
:param nodes: List of nodes owned by this pipeline.
|
||||
:type nodes: list[Node] | None
|
||||
:param pipelines: Child pipelines to compose into this pipeline.
|
||||
:type pipelines: list[Pipeline] | None
|
||||
:param inputs: External input messages (not produced by any node).
|
||||
:type inputs: list[Message] | None
|
||||
:param outputs: Output messages (for documentation purposes).
|
||||
:type outputs: list[Message] | None
|
||||
:param max_workers: Maximum thread pool workers for parallel execution.
|
||||
:type max_workers: int | None
|
||||
:param state: State backend for storing pipeline values.
|
||||
:type state: StateBackend | None
|
||||
|
||||
:example:
|
||||
>>> pipeline = Pipeline(
|
||||
... nodes=[node_a, node_b],
|
||||
... pipelines=[child_pipeline],
|
||||
... max_workers=4
|
||||
... )
|
||||
>>> result = pipeline.run()
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
"_own_nodes",
|
||||
"_child_pipelines",
|
||||
"_parent",
|
||||
"_state",
|
||||
"produces",
|
||||
"dependencies",
|
||||
"_edges",
|
||||
"_execution_order",
|
||||
"_downstream_cache",
|
||||
"_max_workers",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
nodes: list[Node] | None = None,
|
||||
pipelines: list[Pipeline] | None = None,
|
||||
inputs: list[Message] | None = None,
|
||||
outputs: list[Message] | None = None,
|
||||
max_workers: int | None = None,
|
||||
state: StateBackend | None = None,
|
||||
) -> None:
|
||||
self._own_nodes = nodes or []
|
||||
self._child_pipelines = pipelines or []
|
||||
self._parent: Pipeline | None = None
|
||||
self._max_workers = max_workers
|
||||
|
||||
# Set parent reference for child pipelines
|
||||
for child in self._child_pipelines:
|
||||
child._parent = self
|
||||
|
||||
# Collect all nodes and build graph
|
||||
all_nodes = self._collect_all_nodes()
|
||||
|
||||
# Use provided state backend or create default MemoryState
|
||||
self._state: StateBackend = state or MemoryState()
|
||||
|
||||
# Build producer map and dependency graph
|
||||
self.produces: dict[str, Node | None] = {
|
||||
**({m.name: None for m in inputs} if inputs else {}),
|
||||
**{msg: node for node in all_nodes for msg in node.provides},
|
||||
}
|
||||
self.dependencies: dict[Node, frozenset[Node]] = {
|
||||
node: frozenset(
|
||||
self.produces[msg] for msg in node.requires if self.produces.get(msg)
|
||||
)
|
||||
for node in all_nodes
|
||||
}
|
||||
|
||||
# Lazy-initialized caches
|
||||
self._edges: dict[Node, set[Node]] | None = None
|
||||
self._execution_order: list[Node] | None = None
|
||||
self._downstream_cache: dict[Node, list[Node]] = {}
|
||||
|
||||
# Bind all nodes to root pipeline
|
||||
for node in all_nodes:
|
||||
node.bind(self._root)
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Properties
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def _root(self) -> Pipeline:
|
||||
"""
|
||||
Get the root pipeline in the hierarchy.
|
||||
|
||||
:returns: The topmost parent pipeline, or self if no parent.
|
||||
:rtype: Pipeline
|
||||
"""
|
||||
return self._parent._root if self._parent else self
|
||||
|
||||
@property
|
||||
def nodes(self) -> list[Node]:
|
||||
"""
|
||||
All nodes in this pipeline and child pipelines.
|
||||
|
||||
:returns: Flattened list of all nodes.
|
||||
:rtype: list[Node]
|
||||
"""
|
||||
return self._collect_all_nodes()
|
||||
|
||||
@property
|
||||
def edges(self) -> dict[Node, set[Node]]:
|
||||
"""
|
||||
Reverse dependency graph mapping producers to consumers.
|
||||
|
||||
Lazily built on first access.
|
||||
|
||||
:returns: Dict mapping each node to its downstream consumers.
|
||||
:rtype: dict[Node, set[Node]]
|
||||
"""
|
||||
if self._edges is None:
|
||||
self._edges = {node: set() for node in self.nodes}
|
||||
for consumer, producers in self.dependencies.items():
|
||||
for producer in producers:
|
||||
self._edges[producer].add(consumer)
|
||||
return self._edges
|
||||
|
||||
@property
|
||||
def state(self) -> StateBackend:
|
||||
"""
|
||||
State backend for storing pipeline values.
|
||||
|
||||
:returns: State backend.
|
||||
:rtype: StateBackend
|
||||
"""
|
||||
return self._state._data
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Node Access
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def _collect_all_nodes(self) -> list[Node]:
|
||||
"""
|
||||
Recursively collect nodes from this pipeline and all children.
|
||||
|
||||
:returns: List of all nodes in the hierarchy.
|
||||
:rtype: list[Node]
|
||||
"""
|
||||
nodes = list(self._own_nodes)
|
||||
for child in self._child_pipelines:
|
||||
nodes.extend(child._collect_all_nodes())
|
||||
return nodes
|
||||
|
||||
def get_node_by_id(self, nid) -> Node | None:
|
||||
"""
|
||||
Find a node by its ID.
|
||||
|
||||
Searches this pipeline's own nodes first, then child pipelines.
|
||||
|
||||
:param nid: The ID of the node to find.
|
||||
:type nid: Any
|
||||
:returns: The node with the given ID, or None if not found.
|
||||
:rtype: Node | None
|
||||
"""
|
||||
if node := next((n for n in self._own_nodes if n.id is nid), None):
|
||||
return node
|
||||
for child in self._child_pipelines:
|
||||
if node := child.get_node_by_id(nid):
|
||||
return node
|
||||
return None
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Graph Algorithms
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def _topological_sort(self) -> list[Node]:
|
||||
"""
|
||||
Compute topological ordering of nodes using Kahn's algorithm.
|
||||
|
||||
Results are cached for subsequent calls.
|
||||
|
||||
:returns: Nodes in topologically sorted order.
|
||||
:rtype: list[Node]
|
||||
:raises ValueError: If a cycle is detected in the graph.
|
||||
"""
|
||||
if self._execution_order is not None:
|
||||
return self._execution_order
|
||||
|
||||
in_degree = {node: len(deps) for node, deps in self.dependencies.items()}
|
||||
queue = deque(n for n, deg in in_degree.items() if deg == 0)
|
||||
result: list[Node] = []
|
||||
|
||||
while queue:
|
||||
node = queue.popleft()
|
||||
result.append(node)
|
||||
for consumer in self.edges[node]:
|
||||
in_degree[consumer] -= 1
|
||||
if in_degree[consumer] == 0:
|
||||
queue.append(consumer)
|
||||
|
||||
if len(result) != len(self.nodes):
|
||||
raise ValueError("Cycle detected in pipeline graph")
|
||||
|
||||
self._execution_order = result
|
||||
return result
|
||||
|
||||
def _get_downstream(self, start: Node) -> list[Node]:
|
||||
"""
|
||||
Get all downstream nodes from a starting node in topological order.
|
||||
|
||||
Results are cached per start node.
|
||||
|
||||
:param start: The node to find downstream nodes from.
|
||||
:type start: Node
|
||||
:returns: Topologically sorted downstream nodes.
|
||||
:rtype: list[Node]
|
||||
"""
|
||||
if start not in self._downstream_cache:
|
||||
reachable: set[Node] = set()
|
||||
queue = deque([start])
|
||||
while queue:
|
||||
for consumer in self.edges[queue.popleft()]:
|
||||
if consumer not in reachable:
|
||||
reachable.add(consumer)
|
||||
queue.append(consumer)
|
||||
self._downstream_cache[start] = [
|
||||
n for n in self._topological_sort() if n in reachable
|
||||
]
|
||||
return self._downstream_cache[start]
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Synchronous Node Support
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def _version_key(self, msg_name: str) -> str:
|
||||
"""Get the state key for a message's version number."""
|
||||
return f"__version__:{msg_name}"
|
||||
|
||||
def _last_seen_key(self, node_name: str, msg_name: str) -> str:
|
||||
"""Get the state key for the version a node last processed."""
|
||||
return f"__last_seen__:{node_name}:{msg_name}"
|
||||
|
||||
def _increment_message_versions(self, outputs: dict) -> None:
|
||||
"""
|
||||
Increment version numbers for all output messages.
|
||||
|
||||
:param outputs: Dict of message names to values.
|
||||
:type outputs: dict
|
||||
"""
|
||||
for msg_name in outputs:
|
||||
self._state.increment(self._version_key(msg_name))
|
||||
|
||||
def _check_synchronous_ready(self, node: Node) -> tuple[bool, dict[str, int]]:
|
||||
"""
|
||||
Check if a synchronous node is ready to execute.
|
||||
|
||||
A synchronous node is ready when ALL its required inputs have a newer
|
||||
version than what the node last processed.
|
||||
|
||||
:param node: The synchronous node to check.
|
||||
:type node: Node
|
||||
:returns: Tuple of (is_ready, current_versions dict).
|
||||
:rtype: tuple[bool, dict[str, int]]
|
||||
"""
|
||||
if not node.requires:
|
||||
return True, {}
|
||||
|
||||
# Build all keys we need to fetch
|
||||
version_keys = [self._version_key(msg) for msg in node.requires]
|
||||
last_seen_keys = [self._last_seen_key(node.name, msg) for msg in node.requires]
|
||||
|
||||
# Fetch all values atomically
|
||||
all_keys = version_keys + last_seen_keys
|
||||
values = self._state.get_multi(all_keys)
|
||||
|
||||
# Check each input
|
||||
current_versions = {}
|
||||
all_newer = True
|
||||
|
||||
for msg_name in node.requires:
|
||||
version_key = self._version_key(msg_name)
|
||||
last_seen_key = self._last_seen_key(node.name, msg_name)
|
||||
|
||||
current_version = values.get(version_key) or 0
|
||||
last_seen_version = values.get(last_seen_key) or 0
|
||||
|
||||
current_versions[msg_name] = current_version
|
||||
|
||||
# For synchronous nodes, version must be:
|
||||
# 1. Greater than 0 (message has been received at least once)
|
||||
# 2. Strictly greater than last seen (message has been updated since last execution)
|
||||
if current_version == 0 or current_version <= last_seen_version:
|
||||
all_newer = False
|
||||
|
||||
return all_newer, current_versions
|
||||
|
||||
def _try_acquire_synchronous_execution(
|
||||
self,
|
||||
node: Node,
|
||||
current_versions: dict[str, int],
|
||||
) -> bool:
|
||||
"""
|
||||
Attempt to acquire exclusive execution rights for a synchronous node.
|
||||
|
||||
Uses compare-and-swap to atomically verify versions haven't changed
|
||||
and update last_seen versions. This prevents race conditions when
|
||||
multiple threads try to execute the same synchronous node.
|
||||
|
||||
:param node: The node attempting to execute.
|
||||
:type node: Node
|
||||
:param current_versions: The versions that were checked.
|
||||
:type current_versions: dict[str, int]
|
||||
:returns: True if execution rights acquired, False otherwise.
|
||||
:rtype: bool
|
||||
"""
|
||||
if not current_versions:
|
||||
return True
|
||||
|
||||
# Build expected values and updates
|
||||
expected = {}
|
||||
updates = {}
|
||||
|
||||
for msg_name, version in current_versions.items():
|
||||
version_key = self._version_key(msg_name)
|
||||
last_seen_key = self._last_seen_key(node.name, msg_name)
|
||||
|
||||
# Expect the version hasn't changed since we checked
|
||||
expected[version_key] = version
|
||||
# Update last_seen to this version
|
||||
updates[last_seen_key] = version
|
||||
|
||||
return self._state.compare_and_swap_multi(expected, updates)
|
||||
|
||||
def _is_node_ready(self, node: Node, state: StateBackend) -> bool:
|
||||
"""
|
||||
Check if a node is ready to execute.
|
||||
|
||||
For non-synchronous nodes: ready when all required inputs exist in state.
|
||||
For synchronous nodes: ready when all inputs are newer than last processed
|
||||
AND we can acquire exclusive execution rights.
|
||||
|
||||
:param node: The node to check.
|
||||
:type node: Node
|
||||
:param state: State backend containing input values.
|
||||
:type state: StateBackend
|
||||
:returns: True if the node should execute.
|
||||
:rtype: bool
|
||||
"""
|
||||
# First check: all required inputs must exist in state
|
||||
with state.lock():
|
||||
for msg_name in node.requires:
|
||||
if msg_name not in state:
|
||||
return False
|
||||
|
||||
# For non-synchronous nodes, that's all we need
|
||||
if not node.synchronous:
|
||||
return True
|
||||
|
||||
# For synchronous nodes, check if all inputs have been updated
|
||||
is_ready, current_versions = self._check_synchronous_ready(node)
|
||||
|
||||
if not is_ready:
|
||||
return False
|
||||
|
||||
# Try to acquire execution rights atomically
|
||||
# This prevents race conditions when multiple triggers happen concurrently
|
||||
return self._try_acquire_synchronous_execution(node, current_versions)
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Execution Core
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def _execute_node(
|
||||
self,
|
||||
node: Node,
|
||||
state: StateBackend,
|
||||
) -> dict | None:
|
||||
"""
|
||||
Execute a single node without triggering downstream propagation.
|
||||
|
||||
:param node: The node to execute.
|
||||
:type node: Node
|
||||
:param state: State backend containing input values.
|
||||
:type state: StateBackend
|
||||
:returns: Node outputs, or None if no outputs.
|
||||
:rtype: dict | None
|
||||
"""
|
||||
node._pipeline = None
|
||||
try:
|
||||
# Gather inputs with locking
|
||||
with state.lock():
|
||||
inputs = {k: state[k] for k in node.requires if k in state}
|
||||
|
||||
result = node(inputs)
|
||||
|
||||
# Update state with outputs and increment versions
|
||||
if result:
|
||||
with state.lock():
|
||||
state.update(result)
|
||||
# Increment version numbers for synchronous node tracking
|
||||
self._increment_message_versions(result)
|
||||
|
||||
return result
|
||||
finally:
|
||||
node._pipeline = self._root
|
||||
|
||||
def _execute_parallel(
|
||||
self,
|
||||
nodes_subset: set[Node] | None,
|
||||
state: StateBackend,
|
||||
check_ready: bool = False,
|
||||
) -> StateBackend:
|
||||
"""
|
||||
Execute nodes in parallel using dynamic scheduling.
|
||||
|
||||
:param nodes_subset: Subset of nodes to execute, or None for all nodes.
|
||||
:type nodes_subset: set[Node] | None
|
||||
:param state: State backend for inputs/outputs.
|
||||
:type state: StateBackend
|
||||
:param check_ready: If True, check state for required inputs before scheduling.
|
||||
:type check_ready: bool
|
||||
:returns: State backend after execution.
|
||||
:rtype: StateBackend
|
||||
"""
|
||||
target_nodes = nodes_subset or set(self.nodes)
|
||||
if not target_nodes:
|
||||
return state
|
||||
|
||||
# Build in-degree map (only counting deps within target set)
|
||||
in_degree = {
|
||||
n: sum(1 for dep in self.dependencies[n] if dep in target_nodes)
|
||||
for n in target_nodes
|
||||
}
|
||||
|
||||
# Track submitted/completed/skipped nodes
|
||||
submitted: set[Node] = set()
|
||||
skipped: set[Node] = set() # Synchronous nodes that weren't ready
|
||||
node_futures: dict[Node, Future] = {}
|
||||
|
||||
def is_ready(n: Node) -> bool:
|
||||
"""Check if node can be scheduled."""
|
||||
if in_degree[n] != 0:
|
||||
return False
|
||||
if check_ready:
|
||||
# Use the enhanced readiness check that handles synchronous nodes
|
||||
return self._is_node_ready(n, state)
|
||||
return True
|
||||
|
||||
def submit_ready(executor: ThreadPoolExecutor) -> None:
|
||||
"""Submit all currently ready nodes."""
|
||||
for n in target_nodes:
|
||||
if n not in submitted and n not in skipped and is_ready(n):
|
||||
submitted.add(n)
|
||||
node_futures[n] = executor.submit(self._execute_node, n, state)
|
||||
elif n not in submitted and n.synchronous and in_degree[n] == 0:
|
||||
# Mark synchronous nodes that weren't ready as skipped
|
||||
# They may become ready on a future trigger
|
||||
skipped.add(n)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=self._max_workers) as executor:
|
||||
submit_ready(executor)
|
||||
|
||||
while node_futures:
|
||||
done, _ = wait(node_futures.values(), return_when="FIRST_COMPLETED")
|
||||
completed = [n for n, f in node_futures.items() if f in done]
|
||||
|
||||
for n in completed:
|
||||
future = node_futures.pop(n)
|
||||
if exc := future.exception():
|
||||
raise exc
|
||||
|
||||
# Only propagate to downstream nodes if this node produced output
|
||||
# If result is None, the node chose not to forward data (e.g., rate limiting)
|
||||
result = future.result()
|
||||
if result is not None:
|
||||
# Update in-degrees and submit newly ready nodes
|
||||
for consumer in self.edges[n]:
|
||||
if consumer in target_nodes:
|
||||
in_degree[consumer] -= 1
|
||||
submit_ready(executor)
|
||||
|
||||
return state
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Public Execution API
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def run(self, inputs: dict | None = None) -> StateBackend:
|
||||
"""
|
||||
Execute the entire pipeline.
|
||||
|
||||
:param inputs: Initial input values to seed the state.
|
||||
:type inputs: dict | None
|
||||
:returns: State backend containing all computed values.
|
||||
:rtype: StateBackend
|
||||
|
||||
:example:
|
||||
>>> result = pipeline.run({"input": 42})
|
||||
"""
|
||||
# Create a fresh state for this run
|
||||
state = MemoryState()
|
||||
if inputs:
|
||||
state.update(inputs)
|
||||
|
||||
return self._execute_parallel(None, state, check_ready=False)
|
||||
|
||||
def trigger(self, node: Node, outputs: dict | None) -> StateBackend:
|
||||
"""
|
||||
Trigger execution after a node runs externally.
|
||||
|
||||
Updates the shared state with node outputs, increments message versions
|
||||
for synchronous node tracking, and executes all downstream nodes whose
|
||||
dependencies are satisfied.
|
||||
|
||||
:param node: The node that was triggered.
|
||||
:type node: Node
|
||||
:param outputs: Outputs produced by the triggered node.
|
||||
:type outputs: dict | None
|
||||
:returns: Updated shared state.
|
||||
:rtype: StateBackend
|
||||
"""
|
||||
root = self._root
|
||||
state = root._state
|
||||
|
||||
if outputs:
|
||||
with state.lock():
|
||||
state.update(outputs)
|
||||
# Increment version numbers for synchronous node tracking
|
||||
root._increment_message_versions(outputs)
|
||||
|
||||
downstream = set(root._get_downstream(node))
|
||||
if not downstream:
|
||||
return state
|
||||
|
||||
return root._execute_parallel(downstream, state, check_ready=True)
|
||||
|
||||
def reset(self) -> None:
|
||||
"""
|
||||
Reset pipeline state for a new execution cycle.
|
||||
|
||||
Clears all computed values from the shared state.
|
||||
"""
|
||||
self._root._state.clear()
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Visualization
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
def visualize(self, filename: str = "pipeline.png") -> None:
|
||||
"""
|
||||
Save a visualization of the pipeline graph.
|
||||
|
||||
:param filename: Output filename for the image.
|
||||
:type filename: str
|
||||
"""
|
||||
G = nx.DiGraph()
|
||||
G.add_nodes_from(self.nodes)
|
||||
G.add_edges_from(
|
||||
(producer, consumer)
|
||||
for consumer, producers in self.dependencies.items()
|
||||
for producer in producers
|
||||
)
|
||||
|
||||
plt.figure(figsize=(10, 6))
|
||||
pos = nx.spring_layout(G, k=0.8, seed=42)
|
||||
|
||||
node_size = 1800
|
||||
node_radius = (node_size**0.5) / 2
|
||||
|
||||
nx.draw_networkx_nodes(
|
||||
G,
|
||||
pos,
|
||||
node_size=node_size,
|
||||
node_color="#59849B",
|
||||
edgecolors="#59849B",
|
||||
linewidths=1.2,
|
||||
alpha=0.95,
|
||||
)
|
||||
nx.draw_networkx_edges(
|
||||
G,
|
||||
pos,
|
||||
arrows=True,
|
||||
arrowstyle="-|>",
|
||||
arrowsize=18,
|
||||
width=2.8,
|
||||
edge_color="#DE8F6E",
|
||||
connectionstyle="arc3,rad=0.05",
|
||||
min_source_margin=node_radius,
|
||||
min_target_margin=node_radius,
|
||||
)
|
||||
nx.draw_networkx_labels(
|
||||
G,
|
||||
pos,
|
||||
font_size=15,
|
||||
font_color="#F5F9E9",
|
||||
font_weight="bold",
|
||||
)
|
||||
|
||||
plt.axis("off")
|
||||
plt.gca().set_facecolor("#333232")
|
||||
plt.gcf().set_facecolor("#333232")
|
||||
plt.tight_layout()
|
||||
plt.savefig(filename, dpi=200)
|
||||
plt.close()
|
||||
@@ -0,0 +1,478 @@
|
||||
"""
|
||||
State management module with pluggable backends.
|
||||
|
||||
This module provides state storage backends for pipeline execution,
|
||||
supporting both in-memory storage and Redis for distributed execution.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pickle
|
||||
from abc import ABC, abstractmethod
|
||||
from contextlib import contextmanager
|
||||
from threading import RLock
|
||||
from typing import Any, Iterator
|
||||
|
||||
import redis
|
||||
|
||||
|
||||
class StateBackend(ABC):
|
||||
"""
|
||||
Abstract base class for state storage backends.
|
||||
|
||||
Provides a dict-like interface for storing and retrieving state values
|
||||
with support for distributed locking.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
"""
|
||||
Get a value by key.
|
||||
|
||||
:param key: The key to retrieve.
|
||||
:type key: str
|
||||
:param default: Default value if key not found.
|
||||
:type default: Any
|
||||
:returns: The stored value or default.
|
||||
:rtype: Any
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def set(self, key: str, value: Any) -> None:
|
||||
"""
|
||||
Set a value by key.
|
||||
|
||||
:param key: The key to set.
|
||||
:type key: str
|
||||
:param value: The value to store.
|
||||
:type value: Any
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def exists(self, key: str) -> bool:
|
||||
"""
|
||||
Check if a key exists.
|
||||
|
||||
:param key: The key to check.
|
||||
:type key: str
|
||||
:returns: True if key exists.
|
||||
:rtype: bool
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def clear(self) -> None:
|
||||
"""Clear all keys in the state."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def keys(self) -> list[str]:
|
||||
"""
|
||||
Get all keys.
|
||||
|
||||
:returns: List of all keys.
|
||||
:rtype: list[str]
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def update(self, mapping: dict[str, Any]) -> None:
|
||||
"""
|
||||
Update multiple keys atomically.
|
||||
|
||||
:param mapping: Dict of key-value pairs to set.
|
||||
:type mapping: dict[str, Any]
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
@contextmanager
|
||||
def lock(self, timeout: float = 10.0):
|
||||
"""
|
||||
Acquire a lock for thread-safe operations.
|
||||
|
||||
:param timeout: Lock timeout in seconds.
|
||||
:type timeout: float
|
||||
:yields: Lock context.
|
||||
"""
|
||||
...
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Atomic operations for synchronous node support
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def increment(self, key: str) -> int:
|
||||
"""
|
||||
Atomically increment a counter and return the new value.
|
||||
|
||||
:param key: The key to increment.
|
||||
:type key: str
|
||||
:returns: The new value after incrementing.
|
||||
:rtype: int
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def get_multi(self, keys: list[str]) -> dict[str, Any]:
|
||||
"""
|
||||
Get multiple values atomically.
|
||||
|
||||
:param keys: List of keys to retrieve.
|
||||
:type keys: list[str]
|
||||
:returns: Dict mapping keys to values (missing keys have None).
|
||||
:rtype: dict[str, Any]
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def compare_and_swap_multi(
|
||||
self,
|
||||
expected: dict[str, Any],
|
||||
updates: dict[str, Any],
|
||||
) -> bool:
|
||||
"""
|
||||
Atomically check expected values and apply updates if all match.
|
||||
|
||||
This is the key primitive for race-condition-free synchronous node execution.
|
||||
If any expected value doesn't match the current value, no updates are applied.
|
||||
|
||||
:param expected: Dict of keys to their expected values.
|
||||
:type expected: dict[str, Any]
|
||||
:param updates: Dict of keys to their new values (applied only if expected matches).
|
||||
:type updates: dict[str, Any]
|
||||
:returns: True if the swap succeeded, False if expected values didn't match.
|
||||
:rtype: bool
|
||||
"""
|
||||
...
|
||||
|
||||
def __contains__(self, key: str) -> bool:
|
||||
return self.exists(key)
|
||||
|
||||
def __getitem__(self, key: str) -> Any:
|
||||
if not self.exists(key):
|
||||
raise KeyError(key)
|
||||
return self.get(key)
|
||||
|
||||
def __setitem__(self, key: str, value: Any) -> None:
|
||||
self.set(key, value)
|
||||
|
||||
def __iter__(self) -> Iterator[str]:
|
||||
return iter(self.keys())
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.keys())
|
||||
|
||||
|
||||
class MemoryState(StateBackend):
|
||||
"""
|
||||
In-memory state backend using a thread-safe dict.
|
||||
|
||||
Suitable for single-process execution and testing.
|
||||
|
||||
:example:
|
||||
>>> state = MemoryState()
|
||||
>>> state["key"] = "value"
|
||||
>>> state["key"]
|
||||
'value'
|
||||
"""
|
||||
|
||||
__slots__ = ("_data", "_lock")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._data: dict[str, Any] = {}
|
||||
self._lock = RLock() # Reentrant lock for nested access
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
with self._lock:
|
||||
return self._data.get(key, default)
|
||||
|
||||
def set(self, key: str, value: Any) -> None:
|
||||
with self._lock:
|
||||
self._data[key] = value
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
with self._lock:
|
||||
return key in self._data
|
||||
|
||||
def clear(self) -> None:
|
||||
with self._lock:
|
||||
self._data.clear()
|
||||
|
||||
def keys(self) -> list[str]:
|
||||
with self._lock:
|
||||
return list(self._data.keys())
|
||||
|
||||
def update(self, mapping: dict[str, Any]) -> None:
|
||||
with self._lock:
|
||||
self._data.update(mapping)
|
||||
|
||||
@contextmanager
|
||||
def lock(self, timeout: float = 10.0):
|
||||
"""Acquire the internal lock."""
|
||||
acquired = self._lock.acquire(timeout=timeout)
|
||||
if not acquired:
|
||||
raise TimeoutError("Failed to acquire lock")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self._lock.release()
|
||||
|
||||
def increment(self, key: str) -> int:
|
||||
"""Atomically increment a counter and return the new value."""
|
||||
with self._lock:
|
||||
value = self._data.get(key, 0)
|
||||
new_value = value + 1
|
||||
self._data[key] = new_value
|
||||
return new_value
|
||||
|
||||
def get_multi(self, keys: list[str]) -> dict[str, Any]:
|
||||
"""Get multiple values atomically."""
|
||||
with self._lock:
|
||||
return {k: self._data.get(k) for k in keys}
|
||||
|
||||
def compare_and_swap_multi(
|
||||
self,
|
||||
expected: dict[str, Any],
|
||||
updates: dict[str, Any],
|
||||
) -> bool:
|
||||
"""Atomically check expected values and apply updates if all match."""
|
||||
with self._lock:
|
||||
# Check all expected values
|
||||
for key, exp_value in expected.items():
|
||||
actual = self._data.get(key)
|
||||
if actual != exp_value:
|
||||
return False
|
||||
|
||||
# All matched, apply updates
|
||||
self._data.update(updates)
|
||||
return True
|
||||
|
||||
|
||||
class RedisState(StateBackend):
|
||||
"""
|
||||
Redis-based state backend for distributed execution.
|
||||
|
||||
Supports automatic serialization using pickle and distributed locking.
|
||||
Uses a namespace prefix to isolate different pipeline executions.
|
||||
|
||||
:param host: Redis host address.
|
||||
:type host: str
|
||||
:param port: Redis port number.
|
||||
:type port: int
|
||||
:param db: Redis database number.
|
||||
:type db: int
|
||||
:param namespace: Key prefix for isolation.
|
||||
:type namespace: str
|
||||
:param password: Redis password (optional).
|
||||
:type password: str | None
|
||||
:param ttl: Time-to-live for keys in seconds (optional).
|
||||
:type ttl: int | None
|
||||
|
||||
:example:
|
||||
>>> state = RedisState(host="localhost", namespace="my_pipeline")
|
||||
>>> state["result"] = {"value": 42}
|
||||
>>> state["result"]
|
||||
{'value': 42}
|
||||
"""
|
||||
|
||||
__slots__ = ("_client", "_namespace", "_ttl", "_lock_name")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "localhost",
|
||||
port: int = 6379,
|
||||
db: int = 0,
|
||||
namespace: str = "pipeline",
|
||||
password: str | None = None,
|
||||
ttl: int | None = None,
|
||||
) -> None:
|
||||
self._client = redis.Redis(
|
||||
host=host,
|
||||
port=port,
|
||||
db=db,
|
||||
password=password,
|
||||
decode_responses=False,
|
||||
)
|
||||
self._namespace = namespace
|
||||
self._ttl = ttl
|
||||
self._lock_name = f"{namespace}:_lock"
|
||||
|
||||
def _key(self, key: str) -> str:
|
||||
"""Generate namespaced key."""
|
||||
return f"{self._namespace}:{key}"
|
||||
|
||||
def _serialize(self, value: Any) -> bytes:
|
||||
"""Serialize value for storage."""
|
||||
return pickle.dumps(value)
|
||||
|
||||
def _deserialize(self, data: bytes | None) -> Any:
|
||||
"""Deserialize value from storage."""
|
||||
return pickle.loads(data) if data else None
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
data = self._client.get(self._key(key))
|
||||
return self._deserialize(data) if data is not None else default
|
||||
|
||||
def set(self, key: str, value: Any) -> None:
|
||||
data = self._serialize(value)
|
||||
if self._ttl:
|
||||
self._client.setex(self._key(key), self._ttl, data)
|
||||
else:
|
||||
self._client.set(self._key(key), data)
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
return bool(self._client.exists(self._key(key)))
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all keys in the namespace."""
|
||||
pattern = f"{self._namespace}:*"
|
||||
cursor = 0
|
||||
while True:
|
||||
cursor, keys = self._client.scan(cursor, match=pattern, count=100)
|
||||
if keys:
|
||||
self._client.delete(*keys)
|
||||
if cursor == 0:
|
||||
break
|
||||
|
||||
def keys(self) -> list[str]:
|
||||
"""Get all keys in the namespace (without prefix)."""
|
||||
pattern = f"{self._namespace}:*"
|
||||
prefix_len = len(self._namespace) + 1
|
||||
result = []
|
||||
cursor = 0
|
||||
while True:
|
||||
cursor, keys = self._client.scan(cursor, match=pattern, count=100)
|
||||
for k in keys:
|
||||
key_str = k.decode("utf-8")
|
||||
# Skip the lock key
|
||||
if not key_str.endswith(":_lock"):
|
||||
result.append(key_str[prefix_len:])
|
||||
if cursor == 0:
|
||||
break
|
||||
return result
|
||||
|
||||
def update(self, mapping: dict[str, Any]) -> None:
|
||||
"""Update multiple keys atomically using a pipeline."""
|
||||
if not mapping:
|
||||
return
|
||||
pipe = self._client.pipeline()
|
||||
for key, value in mapping.items():
|
||||
data = self._serialize(value)
|
||||
if self._ttl:
|
||||
pipe.setex(self._key(key), self._ttl, data)
|
||||
else:
|
||||
pipe.set(self._key(key), data)
|
||||
pipe.execute()
|
||||
|
||||
@contextmanager
|
||||
def lock(self, timeout: float = 10.0):
|
||||
"""
|
||||
Acquire a distributed lock using Redis.
|
||||
|
||||
:param timeout: Lock timeout in seconds.
|
||||
:type timeout: float
|
||||
:yields: Lock context.
|
||||
:raises TimeoutError: If lock cannot be acquired.
|
||||
"""
|
||||
redis_lock = self._client.lock(
|
||||
self._lock_name,
|
||||
timeout=timeout,
|
||||
blocking_timeout=timeout,
|
||||
)
|
||||
acquired = redis_lock.acquire()
|
||||
if not acquired:
|
||||
raise TimeoutError("Failed to acquire Redis lock")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
try:
|
||||
redis_lock.release()
|
||||
except redis.exceptions.LockError:
|
||||
pass # Lock may have expired
|
||||
|
||||
def ping(self) -> bool:
|
||||
"""
|
||||
Check if Redis connection is available.
|
||||
|
||||
:returns: True if connected.
|
||||
:rtype: bool
|
||||
"""
|
||||
try:
|
||||
return self._client.ping()
|
||||
except redis.ConnectionError:
|
||||
return False
|
||||
|
||||
def increment(self, key: str) -> int:
|
||||
"""Atomically increment a counter using Redis INCR."""
|
||||
return self._client.incr(self._key(key))
|
||||
|
||||
def get_multi(self, keys: list[str]) -> dict[str, Any]:
|
||||
"""Get multiple values atomically using Redis MGET."""
|
||||
if not keys:
|
||||
return {}
|
||||
|
||||
full_keys = [self._key(k) for k in keys]
|
||||
values = self._client.mget(full_keys)
|
||||
|
||||
result = {}
|
||||
for key, value in zip(keys, values):
|
||||
result[key] = self._deserialize(value) if value is not None else None
|
||||
|
||||
return result
|
||||
|
||||
def compare_and_swap_multi(
|
||||
self,
|
||||
expected: dict[str, Any],
|
||||
updates: dict[str, Any],
|
||||
) -> bool:
|
||||
"""
|
||||
Atomically check expected values and apply updates using Redis WATCH/MULTI.
|
||||
|
||||
Uses optimistic locking with WATCH to detect concurrent modifications.
|
||||
If any watched key changes before EXEC, the transaction is aborted.
|
||||
|
||||
:param expected: Dict of keys to their expected values.
|
||||
:type expected: dict[str, Any]
|
||||
:param updates: Dict of keys to their new values.
|
||||
:type updates: dict[str, Any]
|
||||
:returns: True if the swap succeeded, False otherwise.
|
||||
:rtype: bool
|
||||
"""
|
||||
if not expected:
|
||||
# No expectations, just set the updates
|
||||
self.update(updates)
|
||||
return True
|
||||
|
||||
watch_keys = [self._key(k) for k in expected.keys()]
|
||||
|
||||
try:
|
||||
# Start watching the keys for changes
|
||||
self._client.watch(*watch_keys)
|
||||
|
||||
# Check current values
|
||||
current = self.get_multi(list(expected.keys()))
|
||||
|
||||
for key, exp_value in expected.items():
|
||||
if current.get(key) != exp_value:
|
||||
self._client.unwatch()
|
||||
return False
|
||||
|
||||
# All values match - execute the transaction
|
||||
pipe = self._client.pipeline(True) # True = use MULTI/EXEC
|
||||
|
||||
for key, value in updates.items():
|
||||
data = self._serialize(value)
|
||||
if self._ttl:
|
||||
pipe.setex(self._key(key), self._ttl, data)
|
||||
else:
|
||||
pipe.set(self._key(key), data)
|
||||
|
||||
pipe.execute()
|
||||
return True
|
||||
|
||||
except redis.WatchError:
|
||||
# Another client modified one of the watched keys
|
||||
return False
|
||||
@@ -0,0 +1,73 @@
|
||||
from pipeline import Pipeline
|
||||
from state import MemoryState, RedisState
|
||||
from dag_generator import (
|
||||
generate_multi_pipeline_dag,
|
||||
dag_to_pipeline_nodes,
|
||||
)
|
||||
import time
|
||||
import numpy as np
|
||||
import threading
|
||||
|
||||
|
||||
# Thread-safe RNG
|
||||
rng_lock = threading.Lock()
|
||||
rng = np.random.default_rng(1000)
|
||||
|
||||
# Parameters for nodes
|
||||
params = {"rng": rng, "rng_lock": rng_lock}
|
||||
|
||||
state = RedisState(
|
||||
host="localhost",
|
||||
port=6379,
|
||||
namespace="fluksio_pipeline",
|
||||
ttl=3600, # Keys expire after 1 hour
|
||||
)
|
||||
print("Using Redis state backend")
|
||||
assert state.ping(), "Can't connect to redis"
|
||||
|
||||
n_pipelines = 30
|
||||
n_nodes_per_pipeline = 100
|
||||
|
||||
print("\n=== Generating multi-pipeline DAG ===")
|
||||
start = time.time()
|
||||
pipeline_graphs, combined_graph = generate_multi_pipeline_dag(
|
||||
num_pipelines=n_pipelines,
|
||||
nodes_per_pipeline=n_nodes_per_pipeline,
|
||||
cross_pipeline_edges=2,
|
||||
edge_probability=0.4,
|
||||
seed=5,
|
||||
)
|
||||
|
||||
# Create pipelines from each sub-graph
|
||||
child_pipelines = []
|
||||
all_trigger_nodes = []
|
||||
|
||||
for i, G in enumerate(pipeline_graphs):
|
||||
nodes_i, _, triggers_i = dag_to_pipeline_nodes(G, params=params, seed=42 + i)
|
||||
child_pipelines.append(Pipeline(nodes=nodes_i))
|
||||
all_trigger_nodes.extend(triggers_i)
|
||||
print(f"Pipeline {i}: {len(nodes_i)} nodes, {len(triggers_i)} triggers")
|
||||
|
||||
# Compose into parent pipeline
|
||||
pipeline = Pipeline(
|
||||
pipelines=child_pipelines,
|
||||
max_workers=4,
|
||||
state=state,
|
||||
)
|
||||
|
||||
print(f"Generated DAG in {time.time() - start:.2f} seconds")
|
||||
|
||||
# Visualize the complete pipeline
|
||||
if n_pipelines * n_nodes_per_pipeline < 1000:
|
||||
pipeline.visualize("composed_pipeline.png")
|
||||
|
||||
# =============================================================================
|
||||
# Run the pipeline by triggering source nodes
|
||||
# =============================================================================
|
||||
print(f"\n=== Triggering {len(all_trigger_nodes)} source nodes ===")
|
||||
for i, trigger_node in enumerate(all_trigger_nodes):
|
||||
print(f"Triggering node {i + 1}/{len(all_trigger_nodes)}: {trigger_node.name}")
|
||||
pipeline.get_node_by_id(trigger_node.id).trigger()
|
||||
print(f"State contains {len(list(pipeline._state.keys()))} values")
|
||||
|
||||
print("\n=== Pipeline execution complete ===")
|
||||
@@ -0,0 +1,194 @@
|
||||
"""
|
||||
Test script for the PipelineController.
|
||||
|
||||
This script demonstrates the controller's ability to:
|
||||
1. Scan a nodes directory and build a pipeline
|
||||
2. Watch for file changes and update the pipeline dynamically
|
||||
3. Handle errors gracefully without breaking the pipeline
|
||||
4. Accept node code from external sources (simulating frontend input)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Add flow directory to path for imports
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
from controller import PipelineController, NodeStatus
|
||||
from state import RedisState, MemoryState
|
||||
|
||||
|
||||
async def main():
|
||||
"""Run the controller test."""
|
||||
|
||||
print("=" * 60)
|
||||
print("Pipeline Controller Test")
|
||||
print("=" * 60)
|
||||
|
||||
# Use Redis if available, otherwise fall back to memory
|
||||
try:
|
||||
state = RedisState(
|
||||
host="localhost",
|
||||
port=6379,
|
||||
namespace="fluksio_controller_test",
|
||||
ttl=3600,
|
||||
)
|
||||
if state.ping():
|
||||
print("✓ Using Redis state backend")
|
||||
else:
|
||||
raise ConnectionError("Redis ping failed")
|
||||
except Exception as e:
|
||||
print(f"⚠ Redis not available ({e}), using MemoryState")
|
||||
state = MemoryState()
|
||||
|
||||
# Create controller pointing to example nodes
|
||||
nodes_dir = Path(__file__).parent / "nodes_example"
|
||||
|
||||
controller = PipelineController(
|
||||
nodes_dir=nodes_dir,
|
||||
state_backend=state,
|
||||
max_workers=4,
|
||||
node_params={"alert_threshold": 70.0}, # Custom param for alert node
|
||||
)
|
||||
|
||||
# Register callbacks to see what's happening
|
||||
controller.on_node_loaded(
|
||||
lambda nid, node: print(
|
||||
f" ✓ Loaded: {nid} ({len(node.requires)} inputs, {len(node.provides)} outputs)"
|
||||
)
|
||||
)
|
||||
controller.on_node_error(
|
||||
lambda nid, err: print(f" ✗ Error in {nid}: {err.split(chr(10))[0]}")
|
||||
)
|
||||
controller.on_node_removed(lambda nid: print(f" ⊘ Removed: {nid}"))
|
||||
controller.on_pipeline_rebuilt(
|
||||
lambda p: print(
|
||||
f" ⟳ Pipeline rebuilt: {len(p.nodes)} nodes, {len(p.dependencies)} dependencies"
|
||||
)
|
||||
)
|
||||
|
||||
# Start the controller
|
||||
print("\n--- Starting Controller ---")
|
||||
await controller.start()
|
||||
|
||||
# Show current state
|
||||
print("\n--- Current State ---")
|
||||
print(f"Active nodes: {[n.name for n in controller.active_nodes]}")
|
||||
print(f"Error nodes: {list(controller.error_nodes.keys())}")
|
||||
|
||||
if controller.pipeline:
|
||||
print(f"\nPipeline dependency graph:")
|
||||
for node, deps in controller.pipeline.dependencies.items():
|
||||
dep_names = [d.name for d in deps] if deps else ["(trigger)"]
|
||||
print(f" {node.name} <- {dep_names}")
|
||||
|
||||
# Test triggering nodes
|
||||
print("\n--- Testing Node Triggers ---")
|
||||
|
||||
if controller.pipeline:
|
||||
# Find trigger nodes (nodes with no dependencies)
|
||||
trigger_nodes = [
|
||||
n
|
||||
for n in controller.active_nodes
|
||||
if not controller.pipeline.dependencies.get(n)
|
||||
]
|
||||
|
||||
print(f"Trigger nodes: {[n.name for n in trigger_nodes]}")
|
||||
|
||||
# Trigger each trigger node
|
||||
for node in trigger_nodes:
|
||||
print(f"\nTriggering: {node.name}")
|
||||
try:
|
||||
node.trigger()
|
||||
except Exception as e:
|
||||
print(f" Error: {e}")
|
||||
|
||||
# Show state after triggers
|
||||
print(f"\nState after triggers:")
|
||||
for key in controller.state_backend.keys():
|
||||
value = controller.state_backend.get(key)
|
||||
print(f" {key}: {value}")
|
||||
|
||||
# Test dynamic node addition
|
||||
print("\n--- Testing Dynamic Node Addition ---")
|
||||
|
||||
new_node_code = '''
|
||||
"""Dynamically added node that logs all sensor data."""
|
||||
|
||||
from nodes import Node
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params):
|
||||
def log_data(params, temperature=None, humidity=None, comfort_index=None, **kwargs):
|
||||
"""Log all available sensor data."""
|
||||
print(f"[data_logger] Logging: T={temperature}, H={humidity}, Comfort={comfort_index}")
|
||||
return {"log_entry": f"T={temperature}, H={humidity}, C={comfort_index}"}
|
||||
|
||||
return Node(
|
||||
f=log_data,
|
||||
requires=[
|
||||
Message(name="temperature", dtype=float),
|
||||
Message(name="humidity", dtype=float),
|
||||
Message(name="comfort_index", dtype=float),
|
||||
],
|
||||
provides=[Message(name="log_entry", dtype=str)],
|
||||
params=params,
|
||||
name="data_logger",
|
||||
)
|
||||
'''
|
||||
|
||||
print("Adding 'data_logger' node from code...")
|
||||
controller.add_node_from_code("data_logger", new_node_code)
|
||||
|
||||
# Wait for file watcher to pick up the change
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
print(f"\nActive nodes after addition: {[n.name for n in controller.active_nodes]}")
|
||||
|
||||
# Test adding a broken node
|
||||
print("\n--- Testing Error Handling ---")
|
||||
|
||||
broken_node_code = '''
|
||||
"""This node has an intentional error."""
|
||||
|
||||
from nodes import Node
|
||||
from util import Message
|
||||
|
||||
|
||||
def create_node(params):
|
||||
# This will cause a NameError
|
||||
return undefined_variable_that_does_not_exist
|
||||
'''
|
||||
|
||||
print("Adding 'broken_node' with intentional error...")
|
||||
controller.add_node_from_code("broken_node", broken_node_code)
|
||||
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
print(
|
||||
f"\nActive nodes (should exclude broken): {[n.name for n in controller.active_nodes]}"
|
||||
)
|
||||
print(f"Error nodes: {list(controller.error_nodes.keys())}")
|
||||
|
||||
# The pipeline should still work with the other nodes
|
||||
if controller.pipeline:
|
||||
print(f"\nPipeline still has {len(controller.pipeline.nodes)} working nodes")
|
||||
|
||||
# Clean up test nodes
|
||||
print("\n--- Cleaning Up ---")
|
||||
controller.remove_node("data_logger")
|
||||
controller.remove_node("broken_node")
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
print(f"Active nodes after cleanup: {[n.name for n in controller.active_nodes]}")
|
||||
|
||||
# Stop the controller
|
||||
await controller.stop()
|
||||
print("\n✓ Controller stopped")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,270 @@
|
||||
"""
|
||||
Test script for HttpNode functionality.
|
||||
|
||||
This script demonstrates both trigger (receiver) and sender modes of HttpNode,
|
||||
including integration with FastAPI and the PipelineController.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Add flow directory to path for imports
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
from fastapi import FastAPI
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from nodes import HttpNode, Node
|
||||
from pipeline import Pipeline
|
||||
from state import MemoryState
|
||||
from util import Message
|
||||
from controller import PipelineController
|
||||
|
||||
|
||||
def test_http_node_modes():
|
||||
"""Test HttpNode mode detection."""
|
||||
print("=" * 60)
|
||||
print("Testing HttpNode Mode Detection")
|
||||
print("=" * 60)
|
||||
|
||||
# Trigger mode: only provides
|
||||
trigger = HttpNode(
|
||||
url="/api/data",
|
||||
method="POST",
|
||||
provides=[Message(name="value", dtype=float)],
|
||||
)
|
||||
print(f"✓ Trigger node mode: {trigger.mode}")
|
||||
assert trigger.mode == HttpNode.Mode.TRIGGER
|
||||
|
||||
# Sender mode: has requires
|
||||
sender = HttpNode(
|
||||
url="https://httpbin.org/post",
|
||||
method="POST",
|
||||
requires=[Message(name="value", dtype=float)],
|
||||
)
|
||||
print(f"✓ Sender node mode: {sender.mode}")
|
||||
assert sender.mode == HttpNode.Mode.SENDER
|
||||
|
||||
# Both requires and provides = sender mode
|
||||
hybrid = HttpNode(
|
||||
url="https://httpbin.org/post",
|
||||
method="POST",
|
||||
requires=[Message(name="input", dtype=float)],
|
||||
provides=[Message(name="output", dtype=float)],
|
||||
)
|
||||
print(f"✓ Hybrid node mode: {hybrid.mode}")
|
||||
assert hybrid.mode == HttpNode.Mode.SENDER
|
||||
|
||||
# Error: neither requires nor provides
|
||||
try:
|
||||
invalid = HttpNode(url="/api/nothing", method="GET")
|
||||
assert False, "Should have raised ValueError"
|
||||
except ValueError as e:
|
||||
print(f"✓ Correctly rejected invalid config: {e}")
|
||||
|
||||
print("\n✓ All mode detection tests passed!\n")
|
||||
|
||||
|
||||
def test_http_trigger_with_fastapi():
|
||||
"""Test HTTP trigger node with FastAPI."""
|
||||
print("=" * 60)
|
||||
print("Testing HTTP Trigger with FastAPI")
|
||||
print("=" * 60)
|
||||
|
||||
# Create FastAPI app
|
||||
app = FastAPI()
|
||||
|
||||
# Create a simple pipeline with HTTP trigger
|
||||
trigger = HttpNode(
|
||||
url="/api/sensors/temperature",
|
||||
method="POST",
|
||||
provides=[
|
||||
Message(name="temperature", dtype=float),
|
||||
Message(name="sensor_id", dtype=str),
|
||||
],
|
||||
name="temp_trigger",
|
||||
)
|
||||
|
||||
# Create a processing node
|
||||
def process_temp(params, temperature=0.0, sensor_id="unknown", **kwargs):
|
||||
celsius = temperature
|
||||
fahrenheit = celsius * 9 / 5 + 32
|
||||
print(f"[processor] Sensor {sensor_id}: {celsius}°C = {fahrenheit}°F")
|
||||
return {"fahrenheit": fahrenheit}
|
||||
|
||||
processor = Node(
|
||||
f=process_temp,
|
||||
requires=[
|
||||
Message(name="temperature", dtype=float),
|
||||
Message(name="sensor_id", dtype=str),
|
||||
],
|
||||
provides=[Message(name="fahrenheit", dtype=float)],
|
||||
params={},
|
||||
name="temp_processor",
|
||||
)
|
||||
|
||||
# Build pipeline
|
||||
state = MemoryState()
|
||||
pipeline = Pipeline(
|
||||
nodes=[trigger, processor],
|
||||
state=state,
|
||||
max_workers=2,
|
||||
)
|
||||
|
||||
# Register the HTTP route
|
||||
trigger.register_route(app)
|
||||
|
||||
# Test with TestClient
|
||||
client = TestClient(app)
|
||||
|
||||
print("\nSending POST request to /api/sensors/temperature...")
|
||||
response = client.post(
|
||||
"/api/sensors/temperature",
|
||||
json={"temperature": 25.5, "sensor_id": "sensor_001"},
|
||||
)
|
||||
|
||||
print(f"Response status: {response.status_code}")
|
||||
print(f"Response body: {response.json()}")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "triggered"
|
||||
assert data["node"] == "temp_trigger"
|
||||
|
||||
# Check pipeline state
|
||||
print(f"\nPipeline state after trigger:")
|
||||
for key in state.keys():
|
||||
print(f" {key}: {state.get(key)}")
|
||||
|
||||
print("\n✓ HTTP trigger test passed!\n")
|
||||
|
||||
|
||||
def test_http_sender():
|
||||
"""Test HTTP sender node."""
|
||||
print("=" * 60)
|
||||
print("Testing HTTP Sender")
|
||||
print("=" * 60)
|
||||
|
||||
# Create a sender node that posts to httpbin
|
||||
sender = HttpNode(
|
||||
url="https://httpbin.org/post",
|
||||
method="POST",
|
||||
requires=[
|
||||
Message(name="temperature", dtype=float),
|
||||
Message(name="humidity", dtype=float),
|
||||
],
|
||||
name="data_sender",
|
||||
params=dict(
|
||||
timeout=10.0,
|
||||
headers={"X-Custom-Header": "fluksio-test"},
|
||||
),
|
||||
)
|
||||
|
||||
print(f"Sender node created: {sender.name}")
|
||||
print(f" URL: {sender.url}")
|
||||
print(f" Method: {sender.method}")
|
||||
print(f" Mode: {sender.mode}")
|
||||
|
||||
# Note: Actual HTTP call would need async context or be mocked
|
||||
# For this test, we just verify the node is configured correctly
|
||||
print("\n✓ HTTP sender configuration test passed!\n")
|
||||
|
||||
|
||||
async def test_controller_with_http_nodes():
|
||||
"""Test PipelineController with HTTP nodes."""
|
||||
print("=" * 60)
|
||||
print("Testing Controller with HTTP Nodes")
|
||||
print("=" * 60)
|
||||
|
||||
import tempfile
|
||||
|
||||
# Create FastAPI app
|
||||
app = FastAPI()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
nodes_dir = Path(tmpdir) / "nodes"
|
||||
nodes_dir.mkdir()
|
||||
|
||||
# Create an HTTP trigger node file
|
||||
trigger_code = """
|
||||
from nodes import HttpNode
|
||||
from util import Message
|
||||
|
||||
def create_node(params):
|
||||
return HttpNode(
|
||||
url="/api/webhook",
|
||||
method="POST",
|
||||
provides=[
|
||||
Message(name="event_type", dtype=str),
|
||||
Message(name="payload", dtype=object),
|
||||
],
|
||||
name="webhook_receiver",
|
||||
)
|
||||
"""
|
||||
(nodes_dir / "webhook.py").write_text(trigger_code)
|
||||
|
||||
# Create controller with FastAPI app
|
||||
controller = PipelineController(
|
||||
nodes_dir=nodes_dir,
|
||||
max_workers=2,
|
||||
fastapi_app=app,
|
||||
)
|
||||
|
||||
# Register callbacks
|
||||
controller.on_node_loaded(
|
||||
lambda nid, node: print(
|
||||
f" ✓ Loaded: {nid} (mode: {getattr(node, 'mode', 'N/A')})"
|
||||
)
|
||||
)
|
||||
|
||||
# Start controller
|
||||
print("\nStarting controller...")
|
||||
await controller.start()
|
||||
|
||||
print(f"\nActive nodes: {[n.name for n in controller.active_nodes]}")
|
||||
print(f"HTTP trigger nodes: {[n.name for n in controller.http_trigger_nodes]}")
|
||||
print(f"HTTP sender nodes: {[n.name for n in controller.http_sender_nodes]}")
|
||||
|
||||
# Check that route was registered
|
||||
route_paths = [r.path for r in app.routes if hasattr(r, "path")]
|
||||
print(f"\nRegistered routes: {route_paths}")
|
||||
|
||||
if "/api/webhook" in route_paths:
|
||||
print("✓ Webhook route registered!")
|
||||
|
||||
# Test the endpoint
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/api/webhook",
|
||||
json={"event_type": "test", "payload": {"key": "value"}},
|
||||
)
|
||||
print(f"Webhook response: {response.json()}")
|
||||
|
||||
# Stop controller
|
||||
await controller.stop()
|
||||
|
||||
print("\n✓ Controller with HTTP nodes test passed!\n")
|
||||
|
||||
|
||||
def main():
|
||||
"""Run all tests."""
|
||||
print("\n" + "=" * 60)
|
||||
print("HttpNode Test Suite")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
# Synchronous tests
|
||||
test_http_node_modes()
|
||||
test_http_trigger_with_fastapi()
|
||||
test_http_sender()
|
||||
|
||||
# Async tests
|
||||
asyncio.run(test_controller_with_http_nodes())
|
||||
|
||||
print("=" * 60)
|
||||
print("All tests completed successfully!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,447 @@
|
||||
"""
|
||||
Test script for InfluxDbNode functionality.
|
||||
|
||||
This script demonstrates write and read operations with InfluxDB.
|
||||
Requires an InfluxDB instance running.
|
||||
|
||||
To run a local InfluxDB instance:
|
||||
docker run -d -p 8086:8086 \
|
||||
-e DOCKER_INFLUXDB_INIT_MODE=setup \
|
||||
-e DOCKER_INFLUXDB_INIT_USERNAME=admin \
|
||||
-e DOCKER_INFLUXDB_INIT_PASSWORD=adminpass \
|
||||
-e DOCKER_INFLUXDB_INIT_ORG=my-org \
|
||||
-e DOCKER_INFLUXDB_INIT_BUCKET=sensors \
|
||||
-e DOCKER_INFLUXDB_INIT_ADMIN_TOKEN=my-super-secret-token \
|
||||
influxdb:2.7
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
# Add flow directory to path for imports
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
from nodes import InfluxDbNode, Node
|
||||
from pipeline import Pipeline
|
||||
from state import MemoryState
|
||||
from util import Message
|
||||
|
||||
|
||||
# Default test configuration
|
||||
TEST_CONFIG = {
|
||||
"url": "http://10.200.200.115:8086",
|
||||
"token": "***REMOVED-INFLUXDB-TOKEN***==",
|
||||
"org": "strobl",
|
||||
"bucket": "test",
|
||||
}
|
||||
|
||||
|
||||
def test_influxdb_node_validation():
|
||||
"""Test InfluxDbNode parameter validation."""
|
||||
print("=" * 60)
|
||||
print("Testing InfluxDbNode Validation")
|
||||
print("=" * 60)
|
||||
|
||||
# Test missing required params
|
||||
try:
|
||||
node = InfluxDbNode(
|
||||
requires=[Message(name="data", dtype=float)],
|
||||
params={"url": "http://localhost:8086"}, # Missing token, org, bucket
|
||||
)
|
||||
assert False, "Should have raised ValueError"
|
||||
except ValueError as e:
|
||||
print(f"✓ Correctly rejected missing params: {e}")
|
||||
|
||||
# Test neither requires nor provides
|
||||
try:
|
||||
node = InfluxDbNode(params=TEST_CONFIG)
|
||||
assert False, "Should have raised ValueError"
|
||||
except ValueError as e:
|
||||
print(f"✓ Correctly rejected empty requires/provides: {e}")
|
||||
|
||||
# Test valid write-only node
|
||||
writer = InfluxDbNode(
|
||||
requires=[Message(name="temperature", dtype=float)],
|
||||
params={
|
||||
**TEST_CONFIG,
|
||||
"writes": {
|
||||
"temperature": {
|
||||
"measurement": "sensor_data",
|
||||
"field": "temp_celsius",
|
||||
"tags": {"location": "room1"},
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
print(f"✓ Created write-only node: {writer.name}")
|
||||
assert len(writer.requires) == 1
|
||||
assert len(writer.provides) == 0
|
||||
|
||||
# Test valid read-only node
|
||||
reader = InfluxDbNode(
|
||||
provides=[Message(name="avg_temperature", dtype=float)],
|
||||
params={
|
||||
**TEST_CONFIG,
|
||||
"queries": {
|
||||
"avg_temperature": {
|
||||
"measurement": "sensor_data",
|
||||
"field": "temp_celsius",
|
||||
"aggregation": "mean",
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
print(f"✓ Created read-only node: {reader.name}")
|
||||
assert len(reader.requires) == 0
|
||||
assert len(reader.provides) == 1
|
||||
|
||||
# Test combined read/write node
|
||||
combined = InfluxDbNode(
|
||||
requires=[Message(name="raw_temp", dtype=float)],
|
||||
provides=[Message(name="avg_temp", dtype=float)],
|
||||
params={
|
||||
**TEST_CONFIG,
|
||||
"writes": {
|
||||
"raw_temp": {
|
||||
"measurement": "temperature",
|
||||
"field": "value",
|
||||
"tags": {"source": "sensor"},
|
||||
}
|
||||
},
|
||||
"queries": {
|
||||
"avg_temp": {
|
||||
"measurement": "temperature",
|
||||
"field": "value",
|
||||
"aggregation": "mean",
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
print(f"✓ Created combined node: {combined.name}")
|
||||
assert len(combined.requires) == 1
|
||||
assert len(combined.provides) == 1
|
||||
|
||||
print("\n✓ All validation tests passed!\n")
|
||||
|
||||
|
||||
def test_flux_query_building():
|
||||
"""Test Flux query string generation."""
|
||||
print("=" * 60)
|
||||
print("Testing Flux Query Building")
|
||||
print("=" * 60)
|
||||
|
||||
node = InfluxDbNode(
|
||||
provides=[Message(name="value", dtype=float)],
|
||||
params={
|
||||
**TEST_CONFIG,
|
||||
"queries": {
|
||||
"value": {
|
||||
"measurement": "temperature",
|
||||
"field": "celsius",
|
||||
"tags": {"location": "room1", "sensor": "dht22"},
|
||||
"range": "-24h",
|
||||
"aggregation": "mean",
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
query = node._build_flux_query(
|
||||
measurement="temperature",
|
||||
field="celsius",
|
||||
tags={"location": "room1", "sensor": "dht22"},
|
||||
time_range="-24h",
|
||||
aggregation="mean",
|
||||
)
|
||||
|
||||
print(f"Generated Flux query:\n{query}\n")
|
||||
|
||||
# Verify query components
|
||||
assert 'from(bucket: "test")' in query
|
||||
assert "range(start: -24h)" in query
|
||||
assert 'r["_measurement"] == "temperature"' in query
|
||||
assert 'r["_field"] == "celsius"' in query
|
||||
assert 'r["location"] == "room1"' in query
|
||||
assert 'r["sensor"] == "dht22"' in query
|
||||
assert "mean()" in query
|
||||
|
||||
print("✓ Query contains all expected components")
|
||||
|
||||
# Test different aggregations
|
||||
for agg in ["last", "first", "max", "min", "sum", "count"]:
|
||||
query = node._build_flux_query(
|
||||
measurement="test",
|
||||
field="value",
|
||||
tags={},
|
||||
time_range="-1h",
|
||||
aggregation=agg,
|
||||
)
|
||||
assert f"{agg}()" in query
|
||||
print(f"✓ Aggregation '{agg}' works")
|
||||
|
||||
print("\n✓ All query building tests passed!\n")
|
||||
|
||||
|
||||
def test_influxdb_write(skip_if_no_server: bool = True):
|
||||
"""Test writing data to InfluxDB."""
|
||||
print("=" * 60)
|
||||
print("Testing InfluxDB Write")
|
||||
print("=" * 60)
|
||||
|
||||
# Create writer node with write configuration
|
||||
writer = InfluxDbNode(
|
||||
requires=[Message(name="temperature", dtype=float)],
|
||||
params={
|
||||
**TEST_CONFIG,
|
||||
"writes": {
|
||||
"temperature": {
|
||||
"measurement": "test_temperature",
|
||||
"field": "value",
|
||||
"tags": {"location": "test_room", "sensor": "test_sensor"},
|
||||
}
|
||||
},
|
||||
},
|
||||
name="test_writer",
|
||||
)
|
||||
|
||||
# Build a simple pipeline
|
||||
state = MemoryState()
|
||||
pipeline = Pipeline(nodes=[writer], state=state, max_workers=1)
|
||||
|
||||
# Test data to write - now just a simple value!
|
||||
test_data = {"temperature": 25.5}
|
||||
|
||||
print(f"Attempting to write: {test_data}")
|
||||
|
||||
try:
|
||||
writer._write_points(test_data)
|
||||
print("✓ Write successful!")
|
||||
except Exception as e:
|
||||
if skip_if_no_server:
|
||||
print(f"⚠ Write failed (server may not be running): {e}")
|
||||
else:
|
||||
raise
|
||||
|
||||
# Test with runtime tags override
|
||||
test_data_with_tags = {
|
||||
"temperature": {"value": 26.0, "tags": {"sensor": "override_sensor"}}
|
||||
}
|
||||
print(f"Attempting to write with runtime tags: {test_data_with_tags}")
|
||||
|
||||
try:
|
||||
writer._write_points(test_data_with_tags)
|
||||
print("✓ Write with runtime tags successful!")
|
||||
except Exception as e:
|
||||
if skip_if_no_server:
|
||||
print(f"⚠ Write failed (server may not be running): {e}")
|
||||
else:
|
||||
raise
|
||||
|
||||
print("\n✓ Write test completed!\n")
|
||||
|
||||
|
||||
def test_influxdb_read(skip_if_no_server: bool = True):
|
||||
"""Test reading data from InfluxDB."""
|
||||
print("=" * 60)
|
||||
print("Testing InfluxDB Read")
|
||||
print("=" * 60)
|
||||
|
||||
# Create reader node
|
||||
reader = InfluxDbNode(
|
||||
provides=[
|
||||
Message(name="last_temp", dtype=float),
|
||||
Message(name="avg_temp", dtype=float),
|
||||
],
|
||||
params={
|
||||
**TEST_CONFIG,
|
||||
"query_range": "-1h",
|
||||
"queries": {
|
||||
"last_temp": {
|
||||
"measurement": "test_temperature",
|
||||
"field": "value",
|
||||
"tags": {"location": "test_room"},
|
||||
"aggregation": "last",
|
||||
},
|
||||
"avg_temp": {
|
||||
"measurement": "test_temperature",
|
||||
"field": "value",
|
||||
"tags": {"location": "test_room"},
|
||||
"aggregation": "mean",
|
||||
},
|
||||
},
|
||||
},
|
||||
name="test_reader",
|
||||
)
|
||||
|
||||
print(f"Attempting to query data...")
|
||||
|
||||
try:
|
||||
results = reader._query_data()
|
||||
print(f"✓ Query successful! Results: {results}")
|
||||
|
||||
if results:
|
||||
for key, value in results.items():
|
||||
print(f" {key}: {value}")
|
||||
else:
|
||||
print(" No data found (this is OK if no data was written)")
|
||||
except Exception as e:
|
||||
if skip_if_no_server:
|
||||
print(f"⚠ Query failed (server may not be running): {e}")
|
||||
else:
|
||||
raise
|
||||
|
||||
print("\n✓ Read test completed!\n")
|
||||
|
||||
|
||||
def test_influxdb_pipeline_integration(skip_if_no_server: bool = True):
|
||||
"""Test InfluxDbNode in a complete pipeline."""
|
||||
print("=" * 60)
|
||||
print("Testing InfluxDB Pipeline Integration")
|
||||
print("=" * 60)
|
||||
|
||||
# Create a source node that generates sensor data (now just a float!)
|
||||
def generate_sensor_data(params, **kwargs):
|
||||
"""Generate test sensor data."""
|
||||
import random
|
||||
|
||||
return {"temperature": 20.0 + random.random() * 10}
|
||||
|
||||
source = Node(
|
||||
f=generate_sensor_data,
|
||||
requires=[],
|
||||
provides=[Message(name="temperature", dtype=float)],
|
||||
params={},
|
||||
name="sensor_source",
|
||||
)
|
||||
|
||||
# Create InfluxDB writer that receives the sensor data
|
||||
writer = InfluxDbNode(
|
||||
requires=[Message(name="temperature", dtype=float)],
|
||||
params={
|
||||
**TEST_CONFIG,
|
||||
"writes": {
|
||||
"temperature": {
|
||||
"measurement": "pipeline_test",
|
||||
"field": "temp_value",
|
||||
"tags": {
|
||||
"source": "test_pipeline",
|
||||
"run_id": str(int(time.time())),
|
||||
},
|
||||
}
|
||||
},
|
||||
},
|
||||
name="influx_writer",
|
||||
)
|
||||
|
||||
# Build pipeline
|
||||
state = MemoryState()
|
||||
pipeline = Pipeline(
|
||||
nodes=[source, writer],
|
||||
state=state,
|
||||
max_workers=2,
|
||||
)
|
||||
|
||||
print(f"Pipeline created with nodes: {[n.name for n in pipeline.nodes]}")
|
||||
print(f"Dependencies: {pipeline.dependencies}")
|
||||
|
||||
try:
|
||||
# Trigger the source node to start the pipeline
|
||||
print("\nTriggering source node...")
|
||||
result = source.inject({}) # Source has no inputs, just generates output
|
||||
|
||||
print(f"Pipeline execution result: {result}")
|
||||
print("✓ Pipeline integration successful!")
|
||||
except Exception as e:
|
||||
if skip_if_no_server:
|
||||
print(f"⚠ Pipeline test failed (server may not be running): {e}")
|
||||
else:
|
||||
raise
|
||||
|
||||
print("\n✓ Pipeline integration test completed!\n")
|
||||
|
||||
|
||||
async def test_influxdb_async_operations(skip_if_no_server: bool = True):
|
||||
"""Test async write and query operations."""
|
||||
print("=" * 60)
|
||||
print("Testing InfluxDB Async Operations")
|
||||
print("=" * 60)
|
||||
|
||||
# Create a combined node with write and query configs
|
||||
node = InfluxDbNode(
|
||||
requires=[Message(name="temperature", dtype=float)],
|
||||
provides=[Message(name="last_value", dtype=float)],
|
||||
params={
|
||||
**TEST_CONFIG,
|
||||
"writes": {
|
||||
"temperature": {
|
||||
"measurement": "async_test",
|
||||
"field": "value",
|
||||
"tags": {"test": "async"},
|
||||
}
|
||||
},
|
||||
"queries": {
|
||||
"last_value": {
|
||||
"measurement": "async_test",
|
||||
"field": "value",
|
||||
"aggregation": "last",
|
||||
}
|
||||
},
|
||||
},
|
||||
name="async_test_node",
|
||||
)
|
||||
|
||||
try:
|
||||
# Test async write - now just a simple value!
|
||||
print("Testing async write...")
|
||||
await node.write_async({"temperature": 42.0})
|
||||
print("✓ Async write successful!")
|
||||
|
||||
# Wait a moment for the write to be visible
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Test async query
|
||||
print("Testing async query...")
|
||||
results = await node.query_async()
|
||||
print(f"✓ Async query successful! Results: {results}")
|
||||
|
||||
except Exception as e:
|
||||
if skip_if_no_server:
|
||||
print(f"⚠ Async test failed (server may not be running): {e}")
|
||||
else:
|
||||
raise
|
||||
|
||||
print("\n✓ Async operations test completed!\n")
|
||||
|
||||
|
||||
def main():
|
||||
"""Run all tests."""
|
||||
print("\n" + "=" * 60)
|
||||
print("InfluxDbNode Test Suite")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
# Synchronous tests (always run)
|
||||
test_influxdb_node_validation()
|
||||
test_flux_query_building()
|
||||
|
||||
# Tests that require InfluxDB server
|
||||
print("\n" + "-" * 60)
|
||||
print("The following tests require InfluxDB on localhost:8086")
|
||||
print("-" * 60 + "\n")
|
||||
|
||||
try:
|
||||
test_influxdb_write(skip_if_no_server=True)
|
||||
test_influxdb_read(skip_if_no_server=True)
|
||||
test_influxdb_pipeline_integration(skip_if_no_server=True)
|
||||
asyncio.run(test_influxdb_async_operations(skip_if_no_server=True))
|
||||
except KeyboardInterrupt:
|
||||
print("\nTests interrupted by user")
|
||||
|
||||
print("=" * 60)
|
||||
print("All tests completed!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,334 @@
|
||||
"""
|
||||
Test script for MqttNode functionality.
|
||||
|
||||
This script demonstrates both subscriber (trigger) and publisher (sender) modes
|
||||
of MqttNode. Requires an MQTT broker running (e.g., mosquitto).
|
||||
|
||||
To run a local mosquitto broker:
|
||||
docker run -it -p 1883:1883 eclipse-mosquitto mosquitto -c /mosquitto-no-auth.conf
|
||||
|
||||
Or install locally:
|
||||
sudo apt install mosquitto mosquitto-clients
|
||||
sudo systemctl start mosquitto
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
# Add flow directory to path for imports
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
from nodes import MqttNode, Node
|
||||
from pipeline import Pipeline
|
||||
from state import MemoryState
|
||||
from util import Message
|
||||
|
||||
host = "127.0.0.1"
|
||||
|
||||
|
||||
def test_mqtt_node_modes():
|
||||
"""Test MqttNode mode detection."""
|
||||
print("=" * 60)
|
||||
print("Testing MqttNode Mode Detection")
|
||||
print("=" * 60)
|
||||
|
||||
# Subscriber mode: only provides
|
||||
subscriber = MqttNode(
|
||||
topic="sensors/temperature",
|
||||
provides=[Message(name="temperature", dtype=float)],
|
||||
params={"broker_host": host},
|
||||
)
|
||||
print(f"✓ Subscriber node mode: {subscriber.mode}")
|
||||
assert subscriber.mode == MqttNode.Mode.SUBSCRIBER
|
||||
|
||||
# Publisher mode: has requires
|
||||
publisher = MqttNode(
|
||||
topic="actuators/hvac",
|
||||
requires=[Message(name="target_temp", dtype=float)],
|
||||
params={"broker_host": host},
|
||||
)
|
||||
print(f"✓ Publisher node mode: {publisher.mode}")
|
||||
assert publisher.mode == MqttNode.Mode.PUBLISHER
|
||||
|
||||
# Both requires and provides = publisher mode
|
||||
hybrid = MqttNode(
|
||||
topic="devices/thermostat",
|
||||
requires=[Message(name="input", dtype=float)],
|
||||
provides=[Message(name="output", dtype=float)],
|
||||
params={"broker_host": host},
|
||||
)
|
||||
print(f"✓ Hybrid node mode: {hybrid.mode}")
|
||||
assert hybrid.mode == MqttNode.Mode.PUBLISHER
|
||||
|
||||
# Error: neither requires nor provides
|
||||
try:
|
||||
invalid = MqttNode(topic="invalid/topic", params={})
|
||||
assert False, "Should have raised ValueError"
|
||||
except ValueError as e:
|
||||
print(f"✓ Correctly rejected invalid config: {e}")
|
||||
|
||||
print("\n✓ All mode detection tests passed!\n")
|
||||
|
||||
|
||||
def test_mqtt_node_config():
|
||||
"""Test MqttNode configuration from params."""
|
||||
print("=" * 60)
|
||||
print("Testing MqttNode Configuration")
|
||||
print("=" * 60)
|
||||
|
||||
node = MqttNode(
|
||||
topic="test/topic",
|
||||
provides=[Message(name="value", dtype=float)],
|
||||
params={
|
||||
"broker_host": "mqtt.example.com",
|
||||
"broker_port": 8883,
|
||||
"username": "user",
|
||||
"password": "secret",
|
||||
"client_id": "test-client",
|
||||
"qos": 2,
|
||||
"retain": True,
|
||||
"keepalive": 120,
|
||||
},
|
||||
)
|
||||
|
||||
assert node.broker_host == "mqtt.example.com"
|
||||
assert node.broker_port == 8883
|
||||
assert node.username == "user"
|
||||
assert node.password == "secret"
|
||||
assert node.client_id == "test-client"
|
||||
assert node.qos == 2
|
||||
assert node.retain is True
|
||||
assert node.keepalive == 120
|
||||
|
||||
print(f"✓ broker_host: {node.broker_host}")
|
||||
print(f"✓ broker_port: {node.broker_port}")
|
||||
print(f"✓ username: {node.username}")
|
||||
print(f"✓ qos: {node.qos}")
|
||||
print(f"✓ retain: {node.retain}")
|
||||
print(f"✓ keepalive: {node.keepalive}")
|
||||
|
||||
print("\n✓ Configuration test passed!\n")
|
||||
|
||||
|
||||
async def test_mqtt_publisher():
|
||||
"""Test MQTT publisher node (requires broker)."""
|
||||
print("=" * 60)
|
||||
print("Testing MQTT Publisher")
|
||||
print("=" * 60)
|
||||
|
||||
publisher = MqttNode(
|
||||
topic="test/fluksio/output",
|
||||
requires=[
|
||||
Message(name="temperature", dtype=float),
|
||||
Message(name="humidity", dtype=float),
|
||||
],
|
||||
params={
|
||||
"broker_host": host,
|
||||
"broker_port": 1883,
|
||||
"qos": 1,
|
||||
},
|
||||
name="test_publisher",
|
||||
)
|
||||
|
||||
print(f"Publisher node created: {publisher.name}")
|
||||
print(f" Topic: {publisher.topic}")
|
||||
print(f" Mode: {publisher.mode}")
|
||||
print(f" Broker: {publisher.broker_host}:{publisher.broker_port}")
|
||||
|
||||
# Try to publish (will fail if no broker is running)
|
||||
try:
|
||||
# Create a simple pipeline to bind the node
|
||||
state = MemoryState()
|
||||
pipeline = Pipeline(nodes=[publisher], state=state, max_workers=1)
|
||||
|
||||
# Directly call the publish method
|
||||
await publisher._publish_message({"temperature": 25.5, "humidity": 60.0})
|
||||
print("✓ Published message successfully!")
|
||||
except Exception as e:
|
||||
print(f"⚠ Could not publish (broker may not be running): {e}")
|
||||
|
||||
print("\n✓ Publisher test completed!\n")
|
||||
|
||||
|
||||
async def test_mqtt_subscriber():
|
||||
"""Test MQTT subscriber node (requires broker)."""
|
||||
print("=" * 60)
|
||||
print("Testing MQTT Subscriber")
|
||||
print("=" * 60)
|
||||
|
||||
# Create a subscriber node
|
||||
subscriber = MqttNode(
|
||||
topic="test/fluksio/input",
|
||||
provides=[
|
||||
Message(name="value", dtype=float),
|
||||
Message(name="unit", dtype=str),
|
||||
],
|
||||
params={
|
||||
"broker_host": host,
|
||||
"broker_port": 1883,
|
||||
},
|
||||
name="test_subscriber",
|
||||
)
|
||||
|
||||
# Create a processing node that will be triggered
|
||||
received_data = []
|
||||
|
||||
def process_data(params, value=0.0, unit="unknown", **kwargs):
|
||||
print(f"[processor] Received: value={value}, unit={unit}")
|
||||
received_data.append({"value": value, "unit": unit})
|
||||
return {"processed": True}
|
||||
|
||||
processor = Node(
|
||||
f=process_data,
|
||||
requires=[
|
||||
Message(name="value", dtype=float),
|
||||
Message(name="unit", dtype=str),
|
||||
],
|
||||
provides=[Message(name="processed", dtype=bool)],
|
||||
params={},
|
||||
name="data_processor",
|
||||
)
|
||||
|
||||
# Build pipeline
|
||||
state = MemoryState()
|
||||
pipeline = Pipeline(
|
||||
nodes=[subscriber, processor],
|
||||
state=state,
|
||||
max_workers=2,
|
||||
)
|
||||
|
||||
print(f"Subscriber node created: {subscriber.name}")
|
||||
print(f" Topic: {subscriber.topic}")
|
||||
print(f" Mode: {subscriber.mode}")
|
||||
|
||||
# Try to start subscription
|
||||
try:
|
||||
await subscriber.start_subscription()
|
||||
print("✓ Subscription started!")
|
||||
|
||||
# Wait a bit for potential messages
|
||||
print(" Listening for 3 seconds...")
|
||||
await asyncio.sleep(3)
|
||||
|
||||
# Stop subscription
|
||||
await subscriber.stop_subscription()
|
||||
print("✓ Subscription stopped!")
|
||||
|
||||
if received_data:
|
||||
print(f"✓ Received {len(received_data)} messages")
|
||||
else:
|
||||
print(" No messages received (publish to test/fluksio/input to test)")
|
||||
|
||||
except Exception as e:
|
||||
print(f"⚠ Could not subscribe (broker may not be running): {e}")
|
||||
|
||||
print("\n✓ Subscriber test completed!\n")
|
||||
|
||||
|
||||
async def test_mqtt_integration():
|
||||
"""Test full MQTT pub/sub integration (requires broker)."""
|
||||
print("=" * 60)
|
||||
print("Testing MQTT Integration (Pub/Sub)")
|
||||
print("=" * 60)
|
||||
|
||||
import aiomqtt
|
||||
|
||||
topic = "test/fluksio/integration"
|
||||
|
||||
# Create subscriber
|
||||
subscriber = MqttNode(
|
||||
topic=topic,
|
||||
provides=[Message(name="sensor_value", dtype=float)],
|
||||
params={"broker_host": host},
|
||||
name="integration_subscriber",
|
||||
)
|
||||
|
||||
# Track received messages
|
||||
received_values = []
|
||||
|
||||
def track_value(params, sensor_value=0.0, **kwargs):
|
||||
print(f"[tracker] Received sensor_value={sensor_value}")
|
||||
received_values.append(sensor_value)
|
||||
return {"tracked": sensor_value}
|
||||
|
||||
tracker = Node(
|
||||
f=track_value,
|
||||
requires=[Message(name="sensor_value", dtype=float)],
|
||||
provides=[Message(name="tracked", dtype=float)],
|
||||
params={},
|
||||
name="value_tracker",
|
||||
)
|
||||
|
||||
# Build pipeline
|
||||
state = MemoryState()
|
||||
pipeline = Pipeline(
|
||||
nodes=[subscriber, tracker],
|
||||
state=state,
|
||||
max_workers=2,
|
||||
)
|
||||
|
||||
try:
|
||||
# Start subscriber
|
||||
await subscriber.start_subscription()
|
||||
print("✓ Subscriber started")
|
||||
|
||||
# Publish some test messages
|
||||
async with aiomqtt.Client(hostname=host) as client:
|
||||
for i in range(3):
|
||||
value = 20.0 + i * 5
|
||||
payload = json.dumps({"sensor_value": value})
|
||||
await client.publish(topic, payload)
|
||||
print(f" Published: {payload}")
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Wait for messages to be processed
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Stop subscriber
|
||||
await subscriber.stop_subscription()
|
||||
print("✓ Subscriber stopped")
|
||||
|
||||
# Check results
|
||||
print(f"\nReceived values: {received_values}")
|
||||
if len(received_values) == 3:
|
||||
print("✓ All messages received and processed!")
|
||||
else:
|
||||
print(f"⚠ Expected 3 messages, got {len(received_values)}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"⚠ Integration test failed (broker may not be running): {e}")
|
||||
|
||||
print("\n✓ Integration test completed!\n")
|
||||
|
||||
|
||||
def main():
|
||||
"""Run all tests."""
|
||||
print("\n" + "=" * 60)
|
||||
print("MqttNode Test Suite")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
# Synchronous tests (always run)
|
||||
test_mqtt_node_modes()
|
||||
test_mqtt_node_config()
|
||||
|
||||
# Async tests (require broker)
|
||||
print("\n" + "-" * 60)
|
||||
print(f"The following tests require an MQTT broker on {host}:1883")
|
||||
print("-" * 60 + "\n")
|
||||
|
||||
try:
|
||||
asyncio.run(test_mqtt_publisher())
|
||||
asyncio.run(test_mqtt_subscriber())
|
||||
asyncio.run(test_mqtt_integration())
|
||||
except KeyboardInterrupt:
|
||||
print("\nTests interrupted by user")
|
||||
|
||||
print("=" * 60)
|
||||
print("All tests completed!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,512 @@
|
||||
"""
|
||||
Test synchronous node execution in pipelines.
|
||||
|
||||
This module tests the "synchronous" flag which ensures a node only executes
|
||||
when ALL its inputs have changed since the last execution.
|
||||
|
||||
Synchronous nodes are useful for:
|
||||
- Aggregation nodes that need all inputs to be "fresh"
|
||||
- Nodes that should only process complete "batches" of data
|
||||
- Avoiding redundant executions when multiple inputs update at different rates
|
||||
"""
|
||||
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
# Add flow directory to path for imports
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
from state import MemoryState, RedisState
|
||||
from pipeline import Pipeline
|
||||
from nodes import Node
|
||||
from util import Message
|
||||
|
||||
|
||||
def test_synchronous_basic():
|
||||
"""Test basic synchronous node behavior."""
|
||||
print("=" * 60)
|
||||
print("Testing Basic Synchronous Node Behavior")
|
||||
print("=" * 60)
|
||||
|
||||
execution_log = []
|
||||
|
||||
def sensor_a(params, **kwargs):
|
||||
value = params.get("value", 1.0)
|
||||
execution_log.append(f"sensor_a -> {value}")
|
||||
return {"temp_a": value}
|
||||
|
||||
def sensor_b(params, **kwargs):
|
||||
value = params.get("value", 2.0)
|
||||
execution_log.append(f"sensor_b -> {value}")
|
||||
return {"temp_b": value}
|
||||
|
||||
def sync_processor(params, temp_a=0, temp_b=0, **kwargs):
|
||||
result = temp_a + temp_b
|
||||
execution_log.append(f"run_sync_processor({temp_a}, {temp_b}) -> {result}")
|
||||
return {"combined": result}
|
||||
|
||||
def async_processor(params, temp_a=0, temp_b=0, **kwargs):
|
||||
result = temp_a * temp_b
|
||||
execution_log.append(f"run_async_processor({temp_a}, {temp_b}) -> {result}")
|
||||
return {"product": result}
|
||||
|
||||
# Create nodes
|
||||
node_a = Node(
|
||||
f=sensor_a,
|
||||
requires=[],
|
||||
provides=[Message(name="temp_a", dtype=float)],
|
||||
params={"value": 10.0},
|
||||
name="sensor_a",
|
||||
)
|
||||
|
||||
node_b = Node(
|
||||
f=sensor_b,
|
||||
requires=[],
|
||||
provides=[Message(name="temp_b", dtype=float)],
|
||||
params={"value": 20.0},
|
||||
name="sensor_b",
|
||||
)
|
||||
|
||||
# Synchronous node - only executes when BOTH inputs change
|
||||
sync_node = Node(
|
||||
f=sync_processor,
|
||||
requires=[
|
||||
Message(name="temp_a", dtype=float),
|
||||
Message(name="temp_b", dtype=float),
|
||||
],
|
||||
provides=[Message(name="combined", dtype=float)],
|
||||
params={"synchronous": True},
|
||||
name="sync_processor",
|
||||
)
|
||||
|
||||
# Non-synchronous node - executes whenever any input is available
|
||||
async_node = Node(
|
||||
f=async_processor,
|
||||
requires=[
|
||||
Message(name="temp_a", dtype=float),
|
||||
Message(name="temp_b", dtype=float),
|
||||
],
|
||||
provides=[Message(name="product", dtype=float)],
|
||||
params={"synchronous": False},
|
||||
name="async_processor",
|
||||
)
|
||||
|
||||
# Build pipeline
|
||||
state = MemoryState()
|
||||
pipeline = Pipeline(
|
||||
nodes=[node_a, node_b, sync_node, async_node],
|
||||
state=state,
|
||||
max_workers=1, # Single worker for predictable ordering
|
||||
)
|
||||
|
||||
print("\n--- Triggering sensor_a (first time) ---")
|
||||
execution_log.clear()
|
||||
node_a.inject({})
|
||||
print(f"Execution log: {execution_log}")
|
||||
print(f"State: temp_a={state.get('temp_a')}, temp_b={state.get('temp_b')}")
|
||||
|
||||
# Neither processor should run yet (temp_b missing)
|
||||
assert "run_async_processor" not in str(
|
||||
execution_log
|
||||
), "async_processor ran too early"
|
||||
assert "run_sync_processor" not in str(
|
||||
execution_log
|
||||
), "sync_processor ran too early"
|
||||
print("✓ Neither processor ran (temp_b not yet available)")
|
||||
|
||||
print("\n--- Triggering sensor_b (first time) ---")
|
||||
execution_log.clear()
|
||||
node_b.inject({})
|
||||
print(f"Execution log: {execution_log}")
|
||||
|
||||
# Both should run now (first time both inputs are available)
|
||||
assert "run_async_processor" in str(
|
||||
execution_log
|
||||
), "async_processor should have run"
|
||||
assert "run_sync_processor" in str(execution_log), "sync_processor should have run"
|
||||
print("✓ Both processors ran (first time both inputs available)")
|
||||
|
||||
print("\n--- Triggering sensor_a (second time) ---")
|
||||
node_a.params["value"] = 15.0
|
||||
execution_log.clear()
|
||||
node_a.inject({})
|
||||
print(f"Execution log: {execution_log}")
|
||||
|
||||
# async_processor SHOULD run (any input change triggers it)
|
||||
# sync_processor should NOT run (only temp_a changed, not temp_b)
|
||||
assert "run_async_processor" in str(
|
||||
execution_log
|
||||
), "async_processor should have run"
|
||||
assert "run_sync_processor" not in str(
|
||||
execution_log
|
||||
), "sync_processor should NOT have run (only temp_a changed)"
|
||||
print("✓ Only async_processor ran (temp_b hasn't changed)")
|
||||
|
||||
print("\n--- Triggering sensor_b (second time) ---")
|
||||
node_b.params["value"] = 25.0
|
||||
execution_log.clear()
|
||||
node_b.inject({})
|
||||
print(f"Execution log: {execution_log}")
|
||||
|
||||
# Both should run now (both inputs have changed since last sync execution)
|
||||
assert "run_async_processor" in str(
|
||||
execution_log
|
||||
), "async_processor should have run"
|
||||
assert "run_sync_processor" in str(
|
||||
execution_log
|
||||
), "sync_processor should have run (both inputs changed)"
|
||||
print("✓ Both processors ran (both inputs have changed)")
|
||||
|
||||
# Verify final state values
|
||||
assert (
|
||||
pipeline.state.get("combined") == 40.0
|
||||
), f"Expected combined=40.0, got {pipeline.state.get('combined')}"
|
||||
assert (
|
||||
pipeline.state.get("product") == 375.0
|
||||
), f"Expected product=375.0, got {pipeline.state.get('product')}"
|
||||
print(
|
||||
f"\n✓ Final state: combined={pipeline.state.get('combined')}, product={pipeline.state.get('product')}"
|
||||
)
|
||||
|
||||
print("\n✓ Basic synchronous test passed!\n")
|
||||
|
||||
|
||||
def test_synchronous_multiple_triggers():
|
||||
"""Test that synchronous nodes handle rapid successive triggers correctly."""
|
||||
print("=" * 60)
|
||||
print("Testing Synchronous Node with Multiple Rapid Triggers")
|
||||
print("=" * 60)
|
||||
|
||||
sync_execution_count = {"count": 0}
|
||||
async_execution_count = {"count": 0}
|
||||
|
||||
def source_x(params, **kwargs):
|
||||
return {"x": params.get("value", 1)}
|
||||
|
||||
def source_y(params, **kwargs):
|
||||
return {"y": params.get("value", 2)}
|
||||
|
||||
def sync_consumer(params, x=0, y=0, **kwargs):
|
||||
sync_execution_count["count"] += 1
|
||||
return {"sync_out": x + y}
|
||||
|
||||
def async_consumer(params, x=0, y=0, **kwargs):
|
||||
async_execution_count["count"] += 1
|
||||
return {"async_out": x * y}
|
||||
|
||||
node_x = Node(
|
||||
f=source_x,
|
||||
requires=[],
|
||||
provides=[Message(name="x", dtype=int)],
|
||||
params={"value": 1},
|
||||
name="source_x",
|
||||
)
|
||||
|
||||
node_y = Node(
|
||||
f=source_y,
|
||||
requires=[],
|
||||
provides=[Message(name="y", dtype=int)],
|
||||
params={"value": 1},
|
||||
name="source_y",
|
||||
)
|
||||
|
||||
sync_node = Node(
|
||||
f=sync_consumer,
|
||||
requires=[Message(name="x", dtype=int), Message(name="y", dtype=int)],
|
||||
provides=[Message(name="sync_out", dtype=int)],
|
||||
params={"synchronous": True},
|
||||
name="sync_consumer",
|
||||
)
|
||||
|
||||
async_node = Node(
|
||||
f=async_consumer,
|
||||
requires=[Message(name="x", dtype=int), Message(name="y", dtype=int)],
|
||||
provides=[Message(name="async_out", dtype=int)],
|
||||
params={"synchronous": False},
|
||||
name="async_consumer",
|
||||
)
|
||||
|
||||
state = MemoryState()
|
||||
pipeline = Pipeline(
|
||||
nodes=[node_x, node_y, sync_node, async_node],
|
||||
state=state,
|
||||
max_workers=1,
|
||||
)
|
||||
|
||||
# Initial trigger to get both inputs
|
||||
node_x.inject({})
|
||||
node_y.inject({})
|
||||
|
||||
initial_sync = sync_execution_count["count"]
|
||||
initial_async = async_execution_count["count"]
|
||||
|
||||
print(f"After initial triggers: sync={initial_sync}, async={initial_async}")
|
||||
assert initial_sync == 1, "sync_consumer should have run once initially"
|
||||
assert initial_async == 1, "async_consumer should have run once initially"
|
||||
|
||||
# Now trigger X multiple times without triggering Y
|
||||
print("\nTriggering source_x 5 times without changing source_y...")
|
||||
for i in range(5):
|
||||
node_x.params["value"] = 10 + i
|
||||
node_x.inject({})
|
||||
|
||||
final_sync = sync_execution_count["count"]
|
||||
final_async = async_execution_count["count"]
|
||||
|
||||
print(f"After 5 more X triggers: sync={final_sync}, async={final_async}")
|
||||
|
||||
# sync_consumer should NOT have run again (Y didn't change)
|
||||
assert (
|
||||
final_sync == initial_sync
|
||||
), f"sync_consumer should still be at {initial_sync}, got {final_sync}"
|
||||
# async_consumer should have run 5 more times
|
||||
assert (
|
||||
final_async == initial_async + 5
|
||||
), f"async_consumer should be at {initial_async + 5}, got {final_async}"
|
||||
|
||||
print("✓ Synchronous node correctly waited for both inputs to change")
|
||||
|
||||
# Now trigger Y once - sync should run
|
||||
print("\nTriggering source_y once...")
|
||||
node_y.params["value"] = 100
|
||||
node_y.inject({})
|
||||
|
||||
after_y_sync = sync_execution_count["count"]
|
||||
after_y_async = async_execution_count["count"]
|
||||
|
||||
print(f"After Y trigger: sync={after_y_sync}, async={after_y_async}")
|
||||
|
||||
assert after_y_sync == initial_sync + 1, f"sync_consumer should have run once more"
|
||||
assert after_y_async == final_async + 1, f"async_consumer should have run once more"
|
||||
|
||||
print("\n✓ Multiple triggers test passed!\n")
|
||||
|
||||
|
||||
def test_synchronous_race_condition():
|
||||
"""Test that synchronous nodes handle concurrent triggers correctly."""
|
||||
print("=" * 60)
|
||||
print("Testing Synchronous Node Race Condition Handling")
|
||||
print("=" * 60)
|
||||
|
||||
execution_count = {"sync": 0, "async": 0}
|
||||
execution_lock = threading.Lock()
|
||||
|
||||
def source_a(params, **kwargs):
|
||||
return {"data_a": params.get("value", 1)}
|
||||
|
||||
def source_b(params, **kwargs):
|
||||
return {"data_b": params.get("value", 2)}
|
||||
|
||||
def sync_consumer(params, data_a=0, data_b=0, **kwargs):
|
||||
with execution_lock:
|
||||
execution_count["sync"] += 1
|
||||
time.sleep(0.01) # Small delay to increase chance of race conditions
|
||||
return {"sync_result": data_a + data_b}
|
||||
|
||||
def async_consumer(params, data_a=0, data_b=0, **kwargs):
|
||||
with execution_lock:
|
||||
execution_count["async"] += 1
|
||||
return {"async_result": data_a * data_b}
|
||||
|
||||
node_a = Node(
|
||||
f=source_a,
|
||||
requires=[],
|
||||
provides=[Message(name="data_a", dtype=int)],
|
||||
params={"value": 10},
|
||||
name="source_a",
|
||||
)
|
||||
|
||||
node_b = Node(
|
||||
f=source_b,
|
||||
requires=[],
|
||||
provides=[Message(name="data_b", dtype=int)],
|
||||
params={"value": 20},
|
||||
name="source_b",
|
||||
)
|
||||
|
||||
sync_node = Node(
|
||||
f=sync_consumer,
|
||||
requires=[Message(name="data_a", dtype=int), Message(name="data_b", dtype=int)],
|
||||
provides=[Message(name="sync_result", dtype=int)],
|
||||
params={"synchronous": True},
|
||||
name="sync_consumer",
|
||||
)
|
||||
|
||||
async_node = Node(
|
||||
f=async_consumer,
|
||||
requires=[Message(name="data_a", dtype=int), Message(name="data_b", dtype=int)],
|
||||
provides=[Message(name="async_result", dtype=int)],
|
||||
params={"synchronous": False},
|
||||
name="async_consumer",
|
||||
)
|
||||
|
||||
state = MemoryState()
|
||||
pipeline = Pipeline(
|
||||
nodes=[node_a, node_b, sync_node, async_node],
|
||||
state=state,
|
||||
max_workers=4, # Multiple workers for concurrency
|
||||
)
|
||||
|
||||
# Trigger both sensors multiple times concurrently
|
||||
num_rounds = 5
|
||||
print(f"\nTriggering both sensors {num_rounds} times concurrently...")
|
||||
|
||||
threads = []
|
||||
for i in range(num_rounds):
|
||||
node_a.params["value"] = 10 + i
|
||||
node_b.params["value"] = 20 + i
|
||||
|
||||
t1 = threading.Thread(target=lambda: node_a.inject({}))
|
||||
t2 = threading.Thread(target=lambda: node_b.inject({}))
|
||||
threads.extend([t1, t2])
|
||||
t1.start()
|
||||
t2.start()
|
||||
|
||||
# Wait for all threads
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
# Give pipeline time to process
|
||||
time.sleep(0.5)
|
||||
|
||||
print(f"\nExecution counts:")
|
||||
print(f" sync_consumer: {execution_count['sync']}")
|
||||
print(f" async_consumer: {execution_count['async']}")
|
||||
|
||||
# sync_consumer should execute at most num_rounds times
|
||||
# (could be less if some triggers happened before the other input updated)
|
||||
assert (
|
||||
execution_count["sync"] <= num_rounds + 1
|
||||
), f"sync_consumer ran too many times: {execution_count['sync']} (expected <= {num_rounds + 1})"
|
||||
|
||||
# async_consumer will run more frequently
|
||||
assert (
|
||||
execution_count["async"] >= execution_count["sync"]
|
||||
), "async_consumer should run at least as often as sync_consumer"
|
||||
|
||||
print(
|
||||
f"\n✓ sync_consumer ran {execution_count['sync']} times (max expected: {num_rounds + 1})"
|
||||
)
|
||||
print(f"✓ async_consumer ran {execution_count['async']} times")
|
||||
print("\n✓ Race condition test passed!\n")
|
||||
|
||||
|
||||
def test_state_backend_atomic_operations():
|
||||
"""Test the atomic operations used for synchronous node support."""
|
||||
print("=" * 60)
|
||||
print("Testing State Backend Atomic Operations")
|
||||
print("=" * 60)
|
||||
|
||||
state = MemoryState()
|
||||
|
||||
# Test increment
|
||||
print("\nTesting increment...")
|
||||
assert state.increment("counter") == 1
|
||||
assert state.increment("counter") == 2
|
||||
assert state.increment("counter") == 3
|
||||
print("✓ increment works correctly")
|
||||
|
||||
# Test get_multi
|
||||
print("\nTesting get_multi...")
|
||||
state.set("a", 1)
|
||||
state.set("b", 2)
|
||||
state.set("c", 3)
|
||||
result = state.get_multi(["a", "b", "c", "missing"])
|
||||
assert result == {"a": 1, "b": 2, "c": 3, "missing": None}
|
||||
print("✓ get_multi works correctly")
|
||||
|
||||
# Test compare_and_swap_multi - success case
|
||||
print("\nTesting compare_and_swap_multi (success)...")
|
||||
success = state.compare_and_swap_multi(
|
||||
expected={"a": 1, "b": 2}, updates={"a": 10, "b": 20, "new_key": 100}
|
||||
)
|
||||
assert success, "compare_and_swap_multi should succeed"
|
||||
assert state.get("a") == 10
|
||||
assert state.get("b") == 20
|
||||
assert state.get("new_key") == 100
|
||||
print("✓ compare_and_swap_multi succeeded and applied updates")
|
||||
|
||||
# Test compare_and_swap_multi - failure case
|
||||
print("\nTesting compare_and_swap_multi (failure)...")
|
||||
success = state.compare_and_swap_multi(
|
||||
expected={"a": 1, "b": 20}, updates={"a": 999, "b": 999} # a is now 10, not 1
|
||||
)
|
||||
assert not success, "compare_and_swap_multi should fail"
|
||||
assert state.get("a") == 10, "a should be unchanged"
|
||||
assert state.get("b") == 20, "b should be unchanged"
|
||||
print("✓ compare_and_swap_multi correctly rejected mismatched expectation")
|
||||
|
||||
print("\n✓ State backend atomic operations test passed!\n")
|
||||
|
||||
|
||||
def test_concurrent_compare_and_swap():
|
||||
"""Test that compare_and_swap_multi handles concurrent access correctly."""
|
||||
print("=" * 60)
|
||||
print("Testing Concurrent Compare-and-Swap")
|
||||
print("=" * 60)
|
||||
|
||||
state = MemoryState()
|
||||
state.set("version", 0)
|
||||
|
||||
success_count = {"count": 0}
|
||||
failure_count = {"count": 0}
|
||||
count_lock = threading.Lock()
|
||||
|
||||
def try_swap(thread_id: int):
|
||||
"""Try to atomically increment the version."""
|
||||
for _ in range(10):
|
||||
current = state.get("version")
|
||||
success = state.compare_and_swap_multi(
|
||||
expected={"version": current}, updates={"version": current + 1}
|
||||
)
|
||||
with count_lock:
|
||||
if success:
|
||||
success_count["count"] += 1
|
||||
else:
|
||||
failure_count["count"] += 1
|
||||
|
||||
# Start multiple threads trying to increment
|
||||
threads = [threading.Thread(target=try_swap, args=(i,)) for i in range(5)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
final_version = state.get("version")
|
||||
total_attempts = success_count["count"] + failure_count["count"]
|
||||
|
||||
print(f"Total attempts: {total_attempts}")
|
||||
print(f"Successful swaps: {success_count['count']}")
|
||||
print(f"Failed swaps (race lost): {failure_count['count']}")
|
||||
print(f"Final version: {final_version}")
|
||||
|
||||
# The final version should equal the number of successful swaps
|
||||
assert (
|
||||
final_version == success_count["count"]
|
||||
), f"Version mismatch: {final_version} != {success_count['count']}"
|
||||
|
||||
print("\n✓ Concurrent compare-and-swap test passed!\n")
|
||||
|
||||
|
||||
def main():
|
||||
"""Run all synchronous node tests."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Synchronous Node Test Suite")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
# Run all tests
|
||||
test_state_backend_atomic_operations()
|
||||
test_concurrent_compare_and_swap()
|
||||
test_synchronous_basic()
|
||||
test_synchronous_multiple_triggers()
|
||||
test_synchronous_race_condition()
|
||||
|
||||
print("=" * 60)
|
||||
print("All synchronous node tests completed!")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,47 @@
|
||||
from pydantic import BaseModel, Field, field_validator, ConfigDict
|
||||
from typing import Any, Tuple, Type
|
||||
import hashlib
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
"""Defines a message type with validation constraints."""
|
||||
|
||||
model_config = ConfigDict(frozen=True) # Immutable for hashability
|
||||
|
||||
name: str
|
||||
dtype: Type = float
|
||||
shape: Tuple[int, ...] = ()
|
||||
vrange: Tuple[float, float] = (0.0, 1.0)
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
return hashlib.sha256(self.name.encode()).hexdigest()[:12]
|
||||
|
||||
def check(self, value: Any) -> None:
|
||||
if not isinstance(value, self.dtype):
|
||||
raise TypeError(
|
||||
f"{self.name}: expected {self.dtype.__name__}, got {type(value).__name__}"
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Message({self.name})"
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash(self.name)
|
||||
|
||||
|
||||
class NodeParams(BaseModel):
|
||||
"""Parameters passed to node functions."""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
params: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
def __getitem__(self, key: str) -> Any:
|
||||
return self.params[key]
|
||||
|
||||
def __contains__(self, key: str) -> bool:
|
||||
return key in self.params
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
return self.params.get(key, default)
|
||||
@@ -0,0 +1,23 @@
|
||||
import logging
|
||||
|
||||
from sqlmodel import Session
|
||||
|
||||
from app.core.db import engine, init_db
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def init() -> None:
|
||||
with Session(engine) as session:
|
||||
init_db(session)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
logger.info("Creating initial data")
|
||||
init()
|
||||
logger.info("Initial data created")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,33 @@
|
||||
import sentry_sdk
|
||||
from fastapi import FastAPI
|
||||
from fastapi.routing import APIRoute
|
||||
from starlette.middleware.cors import CORSMiddleware
|
||||
|
||||
from app.api.main import api_router
|
||||
from app.core.config import settings
|
||||
|
||||
|
||||
def custom_generate_unique_id(route: APIRoute) -> str:
|
||||
return f"{route.tags[0]}-{route.name}"
|
||||
|
||||
|
||||
if settings.SENTRY_DSN and settings.ENVIRONMENT != "local":
|
||||
sentry_sdk.init(dsn=str(settings.SENTRY_DSN), enable_tracing=True)
|
||||
|
||||
app = FastAPI(
|
||||
title=settings.PROJECT_NAME,
|
||||
openapi_url=f"{settings.API_V1_STR}/openapi.json",
|
||||
generate_unique_id_function=custom_generate_unique_id,
|
||||
)
|
||||
|
||||
# Set all CORS enabled origins
|
||||
if settings.all_cors_origins:
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=settings.all_cors_origins,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
app.include_router(api_router, prefix=settings.API_V1_STR)
|
||||
@@ -0,0 +1,129 @@
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from pydantic import EmailStr
|
||||
from sqlalchemy import DateTime
|
||||
from sqlmodel import Field, Relationship, SQLModel
|
||||
|
||||
|
||||
def get_datetime_utc() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
# Shared properties
|
||||
class UserBase(SQLModel):
|
||||
email: EmailStr = Field(unique=True, index=True, max_length=255)
|
||||
is_active: bool = True
|
||||
is_superuser: bool = False
|
||||
full_name: str | None = Field(default=None, max_length=255)
|
||||
|
||||
|
||||
# Properties to receive via API on creation
|
||||
class UserCreate(UserBase):
|
||||
password: str = Field(min_length=8, max_length=128)
|
||||
|
||||
|
||||
class UserRegister(SQLModel):
|
||||
email: EmailStr = Field(max_length=255)
|
||||
password: str = Field(min_length=8, max_length=128)
|
||||
full_name: str | None = Field(default=None, max_length=255)
|
||||
|
||||
|
||||
# Properties to receive via API on update, all are optional
|
||||
class UserUpdate(UserBase):
|
||||
email: EmailStr | None = Field(default=None, max_length=255) # type: ignore
|
||||
password: str | None = Field(default=None, min_length=8, max_length=128)
|
||||
|
||||
|
||||
class UserUpdateMe(SQLModel):
|
||||
full_name: str | None = Field(default=None, max_length=255)
|
||||
email: EmailStr | None = Field(default=None, max_length=255)
|
||||
|
||||
|
||||
class UpdatePassword(SQLModel):
|
||||
current_password: str = Field(min_length=8, max_length=128)
|
||||
new_password: str = Field(min_length=8, max_length=128)
|
||||
|
||||
|
||||
# Database model, database table inferred from class name
|
||||
class User(UserBase, table=True):
|
||||
id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
|
||||
hashed_password: str
|
||||
created_at: datetime | None = Field(
|
||||
default_factory=get_datetime_utc,
|
||||
sa_type=DateTime(timezone=True), # type: ignore
|
||||
)
|
||||
items: list["Item"] = Relationship(back_populates="owner", cascade_delete=True)
|
||||
|
||||
|
||||
# Properties to return via API, id is always required
|
||||
class UserPublic(UserBase):
|
||||
id: uuid.UUID
|
||||
created_at: datetime | None = None
|
||||
|
||||
|
||||
class UsersPublic(SQLModel):
|
||||
data: list[UserPublic]
|
||||
count: int
|
||||
|
||||
|
||||
# Shared properties
|
||||
class ItemBase(SQLModel):
|
||||
title: str = Field(min_length=1, max_length=255)
|
||||
description: str | None = Field(default=None, max_length=255)
|
||||
|
||||
|
||||
# Properties to receive on item creation
|
||||
class ItemCreate(ItemBase):
|
||||
pass
|
||||
|
||||
|
||||
# Properties to receive on item update
|
||||
class ItemUpdate(ItemBase):
|
||||
title: str | None = Field(default=None, min_length=1, max_length=255) # type: ignore
|
||||
|
||||
|
||||
# Database model, database table inferred from class name
|
||||
class Item(ItemBase, table=True):
|
||||
id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
|
||||
created_at: datetime | None = Field(
|
||||
default_factory=get_datetime_utc,
|
||||
sa_type=DateTime(timezone=True), # type: ignore
|
||||
)
|
||||
owner_id: uuid.UUID = Field(
|
||||
foreign_key="user.id", nullable=False, ondelete="CASCADE"
|
||||
)
|
||||
owner: User | None = Relationship(back_populates="items")
|
||||
|
||||
|
||||
# Properties to return via API, id is always required
|
||||
class ItemPublic(ItemBase):
|
||||
id: uuid.UUID
|
||||
owner_id: uuid.UUID
|
||||
created_at: datetime | None = None
|
||||
|
||||
|
||||
class ItemsPublic(SQLModel):
|
||||
data: list[ItemPublic]
|
||||
count: int
|
||||
|
||||
|
||||
# Generic message
|
||||
class Message(SQLModel):
|
||||
message: str
|
||||
|
||||
|
||||
# JSON payload containing access token
|
||||
class Token(SQLModel):
|
||||
access_token: str
|
||||
token_type: str = "bearer"
|
||||
|
||||
|
||||
# Contents of JWT token
|
||||
class TokenPayload(SQLModel):
|
||||
sub: str | None = None
|
||||
|
||||
|
||||
class NewPassword(SQLModel):
|
||||
token: str
|
||||
new_password: str = Field(min_length=8, max_length=128)
|
||||
@@ -0,0 +1,39 @@
|
||||
import logging
|
||||
|
||||
from sqlalchemy import Engine
|
||||
from sqlmodel import Session, select
|
||||
from tenacity import after_log, before_log, retry, stop_after_attempt, wait_fixed
|
||||
|
||||
from app.core.db import engine
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
max_tries = 60 * 5 # 5 minutes
|
||||
wait_seconds = 1
|
||||
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(max_tries),
|
||||
wait=wait_fixed(wait_seconds),
|
||||
before=before_log(logger, logging.INFO),
|
||||
after=after_log(logger, logging.WARN),
|
||||
)
|
||||
def init(db_engine: Engine) -> None:
|
||||
try:
|
||||
# Try to create session to check if DB is awake
|
||||
with Session(db_engine) as session:
|
||||
session.exec(select(1))
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
raise e
|
||||
|
||||
|
||||
def main() -> None:
|
||||
logger.info("Initializing service")
|
||||
init(engine)
|
||||
logger.info("Service finished initializing")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,123 @@
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import emails # type: ignore
|
||||
import jwt
|
||||
from jinja2 import Template
|
||||
from jwt.exceptions import InvalidTokenError
|
||||
|
||||
from app.core import security
|
||||
from app.core.config import settings
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmailData:
|
||||
html_content: str
|
||||
subject: str
|
||||
|
||||
|
||||
def render_email_template(*, template_name: str, context: dict[str, Any]) -> str:
|
||||
template_str = (
|
||||
Path(__file__).parent / "email-templates" / "build" / template_name
|
||||
).read_text()
|
||||
html_content = Template(template_str).render(context)
|
||||
return html_content
|
||||
|
||||
|
||||
def send_email(
|
||||
*,
|
||||
email_to: str,
|
||||
subject: str = "",
|
||||
html_content: str = "",
|
||||
) -> None:
|
||||
assert settings.emails_enabled, "no provided configuration for email variables"
|
||||
message = emails.Message(
|
||||
subject=subject,
|
||||
html=html_content,
|
||||
mail_from=(settings.EMAILS_FROM_NAME, settings.EMAILS_FROM_EMAIL),
|
||||
)
|
||||
smtp_options = {"host": settings.SMTP_HOST, "port": settings.SMTP_PORT}
|
||||
if settings.SMTP_TLS:
|
||||
smtp_options["tls"] = True
|
||||
elif settings.SMTP_SSL:
|
||||
smtp_options["ssl"] = True
|
||||
if settings.SMTP_USER:
|
||||
smtp_options["user"] = settings.SMTP_USER
|
||||
if settings.SMTP_PASSWORD:
|
||||
smtp_options["password"] = settings.SMTP_PASSWORD
|
||||
response = message.send(to=email_to, smtp=smtp_options)
|
||||
logger.info(f"send email result: {response}")
|
||||
|
||||
|
||||
def generate_test_email(email_to: str) -> EmailData:
|
||||
project_name = settings.PROJECT_NAME
|
||||
subject = f"{project_name} - Test email"
|
||||
html_content = render_email_template(
|
||||
template_name="test_email.html",
|
||||
context={"project_name": settings.PROJECT_NAME, "email": email_to},
|
||||
)
|
||||
return EmailData(html_content=html_content, subject=subject)
|
||||
|
||||
|
||||
def generate_reset_password_email(email_to: str, email: str, token: str) -> EmailData:
|
||||
project_name = settings.PROJECT_NAME
|
||||
subject = f"{project_name} - Password recovery for user {email}"
|
||||
link = f"{settings.FRONTEND_HOST}/reset-password?token={token}"
|
||||
html_content = render_email_template(
|
||||
template_name="reset_password.html",
|
||||
context={
|
||||
"project_name": settings.PROJECT_NAME,
|
||||
"username": email,
|
||||
"email": email_to,
|
||||
"valid_hours": settings.EMAIL_RESET_TOKEN_EXPIRE_HOURS,
|
||||
"link": link,
|
||||
},
|
||||
)
|
||||
return EmailData(html_content=html_content, subject=subject)
|
||||
|
||||
|
||||
def generate_new_account_email(
|
||||
email_to: str, username: str, password: str
|
||||
) -> EmailData:
|
||||
project_name = settings.PROJECT_NAME
|
||||
subject = f"{project_name} - New account for user {username}"
|
||||
html_content = render_email_template(
|
||||
template_name="new_account.html",
|
||||
context={
|
||||
"project_name": settings.PROJECT_NAME,
|
||||
"username": username,
|
||||
"password": password,
|
||||
"email": email_to,
|
||||
"link": settings.FRONTEND_HOST,
|
||||
},
|
||||
)
|
||||
return EmailData(html_content=html_content, subject=subject)
|
||||
|
||||
|
||||
def generate_password_reset_token(email: str) -> str:
|
||||
delta = timedelta(hours=settings.EMAIL_RESET_TOKEN_EXPIRE_HOURS)
|
||||
now = datetime.now(timezone.utc)
|
||||
expires = now + delta
|
||||
exp = expires.timestamp()
|
||||
encoded_jwt = jwt.encode(
|
||||
{"exp": exp, "nbf": now, "sub": email},
|
||||
settings.SECRET_KEY,
|
||||
algorithm=security.ALGORITHM,
|
||||
)
|
||||
return encoded_jwt
|
||||
|
||||
|
||||
def verify_password_reset_token(token: str) -> str | None:
|
||||
try:
|
||||
decoded_token = jwt.decode(
|
||||
token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
|
||||
)
|
||||
return str(decoded_token["sub"])
|
||||
except InvalidTokenError:
|
||||
return None
|
||||
Reference in New Issue
Block a user