dataflow-core 2.1.15rc1__py3-none-any.whl → 2.1.15rc3__py3-none-any.whl

dataflow/models/org_associations.py

@@ -0,0 +1,38 @@
+from sqlalchemy import Column , Integer, String, Boolean, ForeignKey, UniqueConstraint, Enum
+from sqlalchemy.orm import relationship
+from dataflow.db import Base
+import enum
+from dataflow.models.environment import EnvType
+
+class OrganizationUser(Base):
+    """
+    Association Table between USER, ROLE, and ORGANIZATION
+    """
+    __tablename__ = "ORGANIZATION_USER"
+    __table_args__ = (UniqueConstraint('org_id', 'user_id', name='uq_org_user'),)
+
+    org_id = Column(Integer, ForeignKey('ORGANIZATION.id', ondelete="CASCADE"), primary_key=True, nullable=False)
+    user_id = Column(Integer, ForeignKey('USER.user_id', ondelete="CASCADE"), primary_key=True, nullable=False)
+    role_id = Column(Integer, ForeignKey('ROLE.id', ondelete="SET NULL"), nullable=False)
+    active_env_short_name = Column(String, nullable=True)
+    active_env_type = Column(Enum(EnvType), nullable=True)
+    active_server_id = Column(Integer, ForeignKey('CUSTOM_SERVER.id', ondelete="SET NULL"))
+    show_server_page = Column(Boolean, default = True, server_default='true')
+    monthly_allocation = Column(Integer, nullable=True, default=0, server_default='0')
+
+    # Relationships
+    user = relationship("User", back_populates="org_user_assocs")
+    role = relationship("Role", back_populates="org_user_assocs")
+    organization = relationship("Organization", back_populates="org_user_assocs")
+
+class OrganizationServer(Base):
+    __tablename__ = "ORGANIZATION_SERVER"
+
+    org_id = Column(Integer, ForeignKey('ORGANIZATION.id', ondelete="CASCADE"), primary_key=True, nullable=False)
+    server_id = Column(Integer, ForeignKey('SERVER_CONFIG.id', ondelete="CASCADE"), primary_key=True, nullable=False)
+
+class OrganizationAppType(Base):
+    __tablename__ = "ORGANIZATION_APP_TYPE"
+
+    org_id = Column(Integer, ForeignKey('ORGANIZATION.id', ondelete="CASCADE"), primary_key=True, nullable=False)
+    app_type_id = Column(Integer, ForeignKey('APP_TYPE.id', ondelete="CASCADE"), primary_key=True, nullable=False)

dataflow/models/organization.py

@@ -0,0 +1,78 @@
+from sqlalchemy import (
+    Column, Integer, String, Enum, DateTime, ForeignKey, Index, text
+)
+import uuid, enum
+from sqlalchemy.dialects.postgresql import JSONB, UUID
+from sqlalchemy.sql import func
+from sqlalchemy.orm import relationship
+from datetime import datetime
+from dataflow.db import Base
+
+
+class Organization(Base):
+    """
+    Organization model for the database.
+    """
+    __tablename__ = "ORGANIZATION"
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    uid = Column(UUID(as_uuid=True), default=uuid.uuid4, nullable=False, unique=True, server_default=text("gen_random_uuid()"))
+    name = Column(String(255), nullable=False, unique=True)
+    invite_code = Column(String(64), nullable=False, unique=True)
+    email_domain = Column(String(255), nullable=False, unique=True)
+    spark_enabled_zones = Column(JSONB, default=func.json([]), server_default=text("'[]'::jsonb"))  # List of zone IDs where Spark is enabled
+    created_at = Column(DateTime, default=datetime.utcnow, server_default=func.now())
+
+    # Association object link
+    org_user_assocs = relationship("OrganizationUser", back_populates="organization", cascade="all, delete-orphan")
+    custom_servers = relationship("CustomServerConfig")
+    onboarding_requests = relationship("UserOnboarding", back_populates="organization", cascade="all, delete-orphan")
+    servers = relationship("ServerConfig", secondary="ORGANIZATION_SERVER", back_populates="organizations")
+    apps = relationship("AppType", secondary="ORGANIZATION_APP_TYPE", back_populates="organizations")
+    roles = relationship("Role", cascade="all, delete-orphan")
+    environments = relationship("Environment", back_populates="organization")
+
+class OnboardingStatus(enum.Enum):
+    pending = 'pending'
+    rejected = 'rejected'
+    accepted = 'accepted'
+
+class OrganizationOnboarding(Base):
+    __tablename__ = 'ORGANIZATION_ONBOARDING'
+    # This prevents an org from having more than one active ('pending' or 'accepted') application
+    # while allowing multiple 'rejected' entries.
+    __table_args__ = (
+        Index(
+            'idx_pending_org_application',
+            'name',
+            unique=True,
+            postgresql_where=Column('status').in_([
+                OnboardingStatus.pending.value,
+                OnboardingStatus.accepted.value
+            ])
+        ),
+    )
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+
+    name = Column(String(255), nullable=False)
+    age = Column(Integer, nullable=True)
+    domain = Column(String(255), nullable=True)
+    no_of_employees = Column(String(50), nullable=True)
+    address = Column(String(500), nullable=True)
+
+    admin_first_name = Column(String(100), nullable=False)
+    admin_last_name = Column(String(100), nullable=True)
+    admin_designation = Column(String(100), nullable=False)
+    admin_email = Column(String(255), nullable=False, unique=True)
+    admin_username = Column(String(100), nullable=False, unique=True)
+    admin_password = Column(String(255), nullable=False)
+
+    discovery_source = Column(String(255), nullable=True)
+    additional_info = Column(String(1000), nullable=True)
+    size_of_data = Column(String(100), nullable=True)
+
+    user_id = Column(Integer, ForeignKey('USER.user_id'), nullable=False)
+    status = Column(Enum(OnboardingStatus), default=OnboardingStatus.pending, nullable=False)
+
+    user = relationship("User", back_populates="organization_onboarding")

dataflow/models/pinned_projects.py

@@ -1,4 +1,4 @@
-from sqlalchemy import Column, Integer, ForeignKey, DateTime, UniqueConstraint
+from sqlalchemy import Column, Integer, ForeignKey, DateTime, UniqueConstraint, func
 from dataflow.db import Base
 from datetime import datetime
 
@@ -8,7 +8,7 @@ class PinnedProject(Base):
     id = Column(Integer, primary_key=True)
     user_id = Column(Integer, ForeignKey('USER.user_id', ondelete="CASCADE"), index=True)
     project_id = Column(Integer, ForeignKey('PROJECT_DETAIL.project_id', ondelete="CASCADE"), index=True)
-    pinned_at = Column(DateTime, default=datetime.utcnow)
+    pinned_at = Column(DateTime, default=datetime.utcnow, server_default=func.now())
 
     __table_args__ = (
         UniqueConstraint("user_id", "project_id", name="uix_user_project"),

dataflow/models/project_details.py

@@ -1,23 +1,26 @@
-from sqlalchemy import Column, String, Enum, DateTime, Integer, func, ForeignKey
+from sqlalchemy import Column, String, Enum, DateTime, Integer, func, ForeignKey, UniqueConstraint
 from sqlalchemy.orm import relationship
 from dataflow.db import Base
 
 class ProjectDetails(Base):
     __tablename__ = "PROJECT_DETAIL"
+    __table_args__ = (UniqueConstraint('org_id', 'slug', name='uq_project_org_slug'),)
     
     project_id = Column(Integer, primary_key=True, autoincrement=True)
     project_name = Column(String, nullable=False)
     git_url = Column(String)
     git_branch = Column(String, nullable=True)
     git_folder = Column(String, nullable=True)
-    type = Column(String, ForeignKey('RUNTIME_APP_TYPE.name', ondelete="CASCADE"), nullable=False)
-    slug = Column(String, nullable=False, unique=True)
+    type = Column(String, ForeignKey('APP_TYPE.name'), nullable=False)
+    slug = Column(String, nullable=False)
     runtime = Column(String, nullable=False)
-    py_env = Column(String, nullable=True)
+    py_env = Column(Integer, nullable=True)
     launch_url = Column(String, nullable=True)  
-    status = Column(Enum("pending", "created" ,"deployed", "stopped", "failed", name="deployment_status"), default="created")
+    status = Column(Enum("pending", "created" ,"deployed", "stopped", "failed", name="deployment_status"), default="created", server_default="created")
     last_deployed = Column(DateTime, nullable=True)
     created_at = Column(DateTime, nullable=False, server_default=func.now())
     created_by = Column(String, nullable=False)
+    org_id = Column(Integer, ForeignKey('ORGANIZATION.id', ondelete='CASCADE'), nullable=False)
+    airflow_config_file = Column(String, nullable=True)
 
-    app_type = relationship("AppType")
+    app_type = relationship("AppType")

dataflow/models/recent_project_studio.py

@@ -12,7 +12,7 @@ class RecentProjectStudio(Base):
     project_name = Column(String, nullable=False)
     project_path = Column(String, nullable=False)
     last_opened_date = Column(DateTime, server_default=func.now(), nullable=False)
-    remember = Column(Boolean, default=False)
+    remember = Column(Boolean, default=False, server_default='false')
 
     __table_args__ = (
         UniqueConstraint(user_name, project_path, app_name, name='user_name_project_path_app_name_unique'),

dataflow/models/role.py

@@ -1,5 +1,5 @@
 """models.py"""
-from sqlalchemy import Column, Integer, String, Enum
+from sqlalchemy import Column, Integer, String, Enum, ForeignKey, UniqueConstraint
 from sqlalchemy.orm import relationship
 from dataflow.db import Base
 import enum
@@ -11,19 +11,25 @@ class BaseRoleField(enum.Enum):
 
 class Role(Base):
     """
-    Table Role
+    Table ROLE
     """
 
     __tablename__='ROLE'
+    __table_args__ = (
+        UniqueConstraint('name', 'org_id', name='uq_role_name_org'),
+    )
 
     id = Column(Integer, primary_key=True, index=True, autoincrement=True, nullable=False)
-    name = Column(String, unique=True, nullable=False)
+    name = Column(String, nullable=False)
+    org_id = Column(Integer, ForeignKey('ORGANIZATION.id'))
     description = Column(String, nullable=True)
-    base_role = Column(Enum(BaseRoleField), nullable=False, default=BaseRoleField.user)
+    base_role = Column(Enum(BaseRoleField), nullable=False, default=BaseRoleField.user, server_default=BaseRoleField.user.value)
 
-    users = relationship("User", back_populates="role_details", cascade="all, delete-orphan")
-    role_server_assocs = relationship("RoleServer", back_populates="role")
+    # Relationships
     role_zone_assocs = relationship("RoleZone", back_populates="role")
+    org_user_assocs = relationship("OrganizationUser", back_populates="role", cascade="all, delete-orphan")
+    organization = relationship("Organization", back_populates="roles")
+    servers = relationship("CustomServerConfig", secondary="ROLE_SERVER", back_populates="roles")
 
     def __repr__(self):
         return f"<Role(id={self.id}, name='{self.name}', base_role='{self.base_role}')>"

dataflow/models/role_server.py

@@ -5,10 +5,7 @@ from dataflow.db import Base
 
 class RoleServer(Base):
     __tablename__ = 'ROLE_SERVER'
-    __table_args__ = (UniqueConstraint('role_id', 'server_id', name='_role_server_uc'),)
+    __table_args__ = (UniqueConstraint('role_id', 'custom_server_id', name='_role_server_uc'),)
 
     role_id = Column(Integer, ForeignKey('ROLE.id', ondelete="CASCADE"), nullable=False, primary_key=True)
-    server_id = Column(Integer, ForeignKey('CUSTOM_SERVER.id', ondelete="CASCADE"), nullable=False, primary_key=True)
-
-    role = relationship("Role", back_populates="role_server_assocs")
-    server = relationship("CustomServerConfig", back_populates="role_server_assocs")
+    custom_server_id = Column(Integer, ForeignKey('CUSTOM_SERVER.id', ondelete="CASCADE"), nullable=False, primary_key=True)

dataflow/models/role_zone.py

@@ -1,14 +1,19 @@
 from typing import Dict, List, Optional
-from sqlalchemy import Column, Integer, ForeignKey, UniqueConstraint, Boolean
+from sqlalchemy import Column, Integer, ForeignKey, UniqueConstraint, Boolean, Index
 from sqlalchemy.orm import relationship
 from dataflow.db import Base
 
 class RoleZone(Base):
     __tablename__ = 'ROLE_ZONE'
-
+    
     role_id = Column(Integer, ForeignKey('ROLE.id', ondelete="CASCADE"), primary_key=True)
     zone_id = Column(Integer, ForeignKey('DATAFLOW_ZONE.id', ondelete="CASCADE"), primary_key=True)
-    is_default = Column(Boolean, default=False, nullable=False)
+    is_default = Column(Boolean, default=False, nullable=False, server_default='false')
+    
+    __table_args__ = (
+        Index('idx_role_runtime_default', 'role_id', unique=True, 
+              postgresql_where=is_default.is_(True)),
+    )
 
     role = relationship("Role", back_populates="role_zone_assocs")
     zone = relationship("DataflowZone", back_populates="role_zone_assocs")

dataflow/models/server_config.py

@@ -1,4 +1,4 @@
-from sqlalchemy import Column, Integer, String, Boolean, Text, ForeignKey
+from sqlalchemy import Column, Integer, String, Boolean, Text, ForeignKey, text
 from sqlalchemy.dialects.postgresql import JSONB 
 from sqlalchemy.sql import func
 from sqlalchemy.orm import relationship
@@ -14,20 +14,24 @@ class ServerConfig(Base):
     ram = Column(String, nullable=False)
     cpu = Column(String, nullable=False)
     gpu = Column(String)
-    default = Column(Boolean, default=False)
-    tags = Column(JSONB, default=func.json([]))
+    default = Column(Boolean, default=False, server_default='false')
+    tags = Column(JSONB, default=func.json([]), server_default=text("'[]'::jsonb"))
     description = Column(Text, nullable=True)
-    kubespawner_override = Column(JSONB, default=func.json({}))
+    kubespawner_override = Column(JSONB, default=func.json({}), server_default=text("'{}'::jsonb"))
 
+    # Relationships
+    organizations = relationship("Organization", secondary="ORGANIZATION_SERVER", back_populates="servers")
 
 class CustomServerConfig(Base):
     __tablename__ = "CUSTOM_SERVER"
 
     id = Column(Integer, primary_key=True, autoincrement=True)
     base_server_id = Column(Integer, ForeignKey(ServerConfig.id, ondelete="CASCADE"), nullable=False)
+    org_id = Column(Integer, ForeignKey('ORGANIZATION.id', ondelete="CASCADE"), nullable=False)
     display_name = Column(String, nullable=False, unique=True, index=True)
     description = Column(Text, nullable=True)
     
     # Relationship to the server_config table
     server_config = relationship(ServerConfig)
-    role_server_assocs = relationship("RoleServer", back_populates="server", cascade="all, delete-orphan")
+    organization = relationship("Organization", back_populates="custom_servers")
+    roles = relationship("Role", secondary="ROLE_SERVER", back_populates="servers")

dataflow/models/team.py

@@ -1,5 +1,5 @@
 """models.py"""
-from sqlalchemy import Column, Integer, String
+from sqlalchemy import Column, Integer, String, ForeignKey, UniqueConstraint
 from sqlalchemy.orm import relationship
 from dataflow.db import Base
 
@@ -9,9 +9,15 @@ class Team(Base):
     """
 
     __tablename__='TEAM'
+    __table_args__ = (
+        UniqueConstraint('team_name', 'org_id', name='uc_team_name_org_id'),
+    )
 
     team_id = Column(Integer, primary_key=True, index=True, autoincrement=True, nullable=False)
-    team_name = Column(String, unique=True, nullable=False)
+    team_name = Column(String, nullable=False)
+    org_id = Column(Integer, ForeignKey('ORGANIZATION.id', ondelete="CASCADE"), nullable=False)
     description = Column(String, nullable=True)
-    
-    user_team_assocs = relationship("UserTeam", back_populates="team")
+
+    # relationships
+    users = relationship("User", secondary="USER_TEAM", back_populates="teams")
+    organization = relationship("Organization")

dataflow/models/user.py

@@ -1,31 +1,68 @@
 """models.py"""
-from sqlalchemy import Column, Integer, String, Boolean, LargeBinary, ForeignKey
+from sqlalchemy import Column, Integer, String, Boolean, LargeBinary, Enum, ForeignKey, DateTime, func
+from sqlalchemy.dialects.postgresql import ENUM
+from sqlalchemy import Index
 from sqlalchemy.orm import relationship
 from dataflow.db import Base
+import enum
 
 class User(Base):
     """
     Table USER
     """
-
-    __tablename__='USER'
+    __tablename__ = 'USER'
 
     user_id = Column(Integer, primary_key=True, index=True, autoincrement=True, nullable=False)
     user_name = Column(String, unique=True, nullable=False)
     first_name = Column(String)
     last_name = Column(String)
     email = Column(String, unique=True)
-    role_id = Column(Integer, ForeignKey('ROLE.id'), nullable=False)
     image = Column(LargeBinary)
     image_url = Column(String, nullable=True)
-    active = Column(Boolean, nullable=False, default=True)
+    active = Column(Boolean, nullable=False, default=True, server_default='true')
     password = Column(String, nullable=False)
-    active_env = Column(String)
-    active_env_type = Column(String, nullable=True)
-    current_server = Column(String)
-    show_server_page = Column(Boolean, default = True)
-    monthly_allocation = Column(Integer, nullable=True, default=0)
+    active_org_id = Column(Integer, ForeignKey('ORGANIZATION.id'))
+
+    # Relationships
+    org_user_assocs = relationship("OrganizationUser", back_populates="user", cascade="all, delete-orphan")
+    teams = relationship("Team", secondary="USER_TEAM", back_populates="users")
+    onboarding_requests = relationship("UserOnboarding", back_populates="user", cascade="all, delete-orphan")
+    organization_onboarding = relationship("OrganizationOnboarding", back_populates="user", cascade="all, delete-orphan")
+
+
+class OnboardingStatus(enum.Enum):
+    pending = 'pending'
+    rejected = 'rejected'
+    accepted = 'accepted'
+
+class UserOnboarding(Base):
+    """
+    SQLAlchemy model for the "USER_ONBOARDING" table.
+    This table stores user applications to organizations.
+    """
+    __tablename__ = "USER_ONBOARDING"
+    # This prevents a user from having more than one active ('pending' or 'accepted') application
+    # for a given organization, while allowing multiple 'rejected' entries.
+    __table_args__ = (
+        Index(
+            'idx_pending_user_org_application',
+            'user_id',
+            'org_id',
+            unique=True,
+            postgresql_where=Column('status').in_([
+                OnboardingStatus.pending.value,
+                OnboardingStatus.accepted.value
+            ])
+        ),
+    )
 
-    role_details = relationship("Role")
+    id = Column(Integer, primary_key=True, autoincrement=True, nullable=False)
+    user_id = Column(Integer, ForeignKey("USER.user_id", ondelete="CASCADE"), nullable=False)
+    org_id = Column(Integer, ForeignKey("ORGANIZATION.id", ondelete="CASCADE"), nullable=False)
+    status = Column(Enum(OnboardingStatus, name='onboarding_status'), nullable=False, default=OnboardingStatus.pending.value, server_default='pending')
+    created_at = Column(DateTime, default=func.now(), nullable=False, server_default=func.now())
+    updated_at = Column(DateTime, default=func.now(), onupdate=func.now(), nullable=False, server_default=func.now())
 
-    user_team_assocs = relationship("UserTeam", back_populates="user")
+    # Relationships
+    user = relationship("User", back_populates="onboarding_requests")
+    organization = relationship("Organization", back_populates="onboarding_requests")

dataflow/models/user_team.py

@@ -8,7 +8,4 @@ class UserTeam(Base):
     __table_args__ = (UniqueConstraint('user_id', 'team_id', name='_user_team_uc'),)
 
     user_id = Column(Integer, ForeignKey('USER.user_id', ondelete="CASCADE"), nullable=False, primary_key=True)
-    team_id = Column(Integer, ForeignKey('TEAM.team_id', ondelete="CASCADE"), nullable=False, primary_key=True)
-
-    user = relationship("User", back_populates="user_team_assocs")
-    team = relationship("Team", back_populates="user_team_assocs")
+    team_id = Column(Integer, ForeignKey('TEAM.team_id', ondelete="CASCADE"), nullable=False, primary_key=True)

dataflow/models/variables.py

@@ -14,7 +14,8 @@ class Variable(Base):
     __tablename__ = 'VARIABLE'
 
     id = Column(Integer, primary_key=True, index=True, autoincrement=True, nullable=False)
-    key = Column(String, nullable=False)
+    key = Column(String, index=True, nullable=False)
+    org_id = Column(Integer, ForeignKey("ORGANIZATION.id"), index=True, nullable=False)
     value = Column(Text, nullable=False)
     type = Column(String, nullable=False)
     description = Column(Text, nullable=True)
@@ -24,11 +25,12 @@ class Variable(Base):
     created_at = Column(DateTime, server_default=func.now())
     updated_at = Column(DateTime, server_default=func.now(), onupdate=func.now())
     created_by = Column(String, ForeignKey('USER.user_name'), nullable=True)
-    is_active = Column(Boolean, default=True, nullable=False)
+    is_active = Column(Boolean, default=True, nullable=False, server_default='true')
     datatype = Column(Enum(DataType, name="data_type"), nullable=False)
-    set_as_env = Column(Boolean, default=False, nullable=False)
+    set_as_env = Column(Boolean, default=False, nullable=False, server_default='false')
+
 
     __table_args__ = (
         CheckConstraint(type.in_(['variable', 'secret']), name='check_variable_type'),
-        UniqueConstraint('key', 'runtime', 'slug', 'created_by', name='unique_key'),
+        UniqueConstraint('key', 'org_id', 'runtime', 'slug', 'created_by', name='unique_key'),
     )

dataflow/secrets_manager/service.py

@@ -117,27 +117,29 @@ class _BaseAccessor(ABC):
 # These classes implement the logic for building the vault path based on the context.
 # --------------------------------------------------------------------------
 class _RuntimeAccessor(_BaseAccessor):
-    def __init__(self, secret_manager: SecretManager, runtime_env: str, slug: str = None):
+    def __init__(self, secret_manager: SecretManager, org_id: str, runtime_env: str, slug: str = None):
         super().__init__(secret_manager)
+        self.org_id = org_id
         self.runtime_env = runtime_env
         self.slug = slug
 
     def _get_vault_path(self, secret_type: str, key: str) -> str:
         # Special case for git-ssh in runtime context
         if secret_type == "git-ssh":
-            return f"{self.runtime_env}-{secret_type}-{self.slug}"
+            return f"{self.org_id}-{self.runtime_env}-{secret_type}-{self.slug}"
 
         # Standard format for all other secret types
         context = self.slug if self.slug else "global"
-        return f"{self.runtime_env}-{context}-{secret_type}-{key}"
+        return f"{self.org_id}-{self.runtime_env}-{context}-{secret_type}-{key}"
 
 class _StudioAccessor(_BaseAccessor):
-    def __init__(self, secret_manager: SecretManager, user_name: str):
+    def __init__(self, secret_manager: SecretManager, org_id: str, user_name: str):
         super().__init__(secret_manager)
+        self.org_id = org_id
         self.user_name = user_name
 
     def _get_vault_path(self, secret_type: str, key: str) -> str:
-        return f"{self.user_name}-{secret_type}-{key}"
+        return f"{self.org_id}-{self.user_name}-{secret_type}-{key}"
 
 # --------------------------------------------------------------------------
 # PUBLIC INTERFACE CLASS
@@ -147,10 +149,10 @@ class SecretsService:
     def __init__(self, secret_manager: SecretManager):
         self._secret_manager = secret_manager
 
-    def runtime(self, env: str, slug: str = None) -> _RuntimeAccessor:
+    def runtime(self, org_id: int, env: str, slug: str = None) -> _RuntimeAccessor:
         """Sets the context to RUNTIME and returns the appropriate accessor."""
-        return _RuntimeAccessor(self._secret_manager, env, slug)
+        return _RuntimeAccessor(self._secret_manager, str(org_id), env, slug)
 
-    def studio(self, user: str) -> _StudioAccessor:
+    def studio(self, org_id: int, user: str) -> _StudioAccessor:
         """Sets the context to STUDIO and returns the appropriate accessor."""
-        return _StudioAccessor(self._secret_manager, user)
+        return _StudioAccessor(self._secret_manager, str(org_id), user)

dataflow/utils/get_current_user.py

@@ -1,37 +1,60 @@
-from fastapi import HTTPException, status
+from fastapi import HTTPException
 from sqlalchemy.orm import Session
-from dataflow.models import user as m_user  
-from dataflow.models import session as m_session
+from sqlalchemy import and_
+from dataflow.models import (
+    user as m_user,
+    session as m_session,
+    org_associations as m_org_associations,
+    role as m_role,
+    organization as m_organization
+)
 
 def get_user_from_session(session_id: str, db: Session):
     """
-    Retrieve a user based on session ID.
-    
-    Args:
-        session_id (str): The unique session identifier
-        db (Session): Database session
-        
-    Returns:
-        User: User object if found
-        
-    Raises:
-        HTTPException: If session is invalid or user not found
+    Retrieve user information based on the session ID.
     """
-    session_record = db.query(m_session.Session).filter(m_session.Session.session_id == session_id).first()
-    if not session_record:
-        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid session")
+    user = (
+        db.query(
+            m_user.User,
+            m_org_associations.OrganizationUser.role_id,
+            m_role.Role.name,
+            m_role.Role.base_role,
+            m_org_associations.OrganizationUser.active_server_id,
+            m_org_associations.OrganizationUser.show_server_page,
+            m_org_associations.OrganizationUser.active_env_short_name,
+            m_org_associations.OrganizationUser.active_env_type,
+            m_org_associations.OrganizationUser.monthly_allocation,
+            m_organization.Organization.uid
+        )
+        .join(m_session.Session, m_session.Session.user_id == m_user.User.user_id)
+        .outerjoin(
+            m_org_associations.OrganizationUser,
+            and_(
+                m_org_associations.OrganizationUser.user_id == m_user.User.user_id,
+                m_org_associations.OrganizationUser.org_id == m_user.User.active_org_id
+            )
+        )
+        .outerjoin(m_role.Role, m_role.Role.id == m_org_associations.OrganizationUser.role_id)
+        .outerjoin(
+            m_organization.Organization,
+            m_organization.Organization.id == m_user.User.active_org_id  # join to Organization
+        )
+        .filter(m_session.Session.session_id == session_id)
+        .first()
+    )
 
-    user = db.query(m_user.User).filter(m_user.User.user_id == session_record.user_id).first()
     if not user:
-        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="User not found")
-
-    base_role = user.role_details.base_role
-    role_id = user.role_details.id
-    role_name = user.role_details.name
-    user.base_role = base_role
-    user.role = role_name
-    user.role_id = role_id
-    
-    return user
+        raise HTTPException(status_code=401, detail="Invalid session")
 
+    user_obj, role_id, role_name, base_role, active_server_id, show_server_page, active_env_short_name, active_env_type, monthly_allocation, uid = user
+    user_obj.role_id = role_id
+    user_obj.role = role_name
+    user_obj.base_role = base_role
+    user_obj.current_server_id = active_server_id
+    user_obj.show_server_page = show_server_page
+    user_obj.active_env = active_env_short_name
+    user_obj.active_env_type = active_env_type
+    user_obj.monthly_allocation = monthly_allocation
+    user_obj.org_uid = str(uid)
 
+    return user_obj

{dataflow_core-2.1.15rc1.dist-info → dataflow_core-2.1.15rc3.dist-info}/METADATA

@@ -1,10 +1,11 @@
 Metadata-Version: 2.4
 Name: dataflow-core
-Version: 2.1.15rc1
+Version: 2.1.15rc3
 Summary: Dataflow core package
 Author: Dataflow
 Author-email: 
 Requires-Dist: sqlalchemy
+Requires-Dist: alembic
 Requires-Dist: boto3
 Requires-Dist: psycopg2-binary
 Requires-Dist: pymysql