dataflow-core 2.1.14rc1__py3-none-any.whl → 2.1.15__py3-none-any.whl

authenticator/dataflowairflowauthenticator.py

@@ -4,7 +4,7 @@ from flask_appbuilder.security.views import expose
 from flask_login import login_user
 from airflow.www.security import FabAirflowSecurityManagerOverride
 from dataflow.dataflow import Dataflow
-import logging, os
+import logging
 
 logging.basicConfig(
     level=logging.INFO,
@@ -29,15 +29,9 @@ class DataflowAuthDBView(AuthDBView):
             
             user_details = dataflow.auth(session_id)
             logger.info(f"User details retrieved for: {user_details['user_name']}")
-            if os.getenv("RUNTIME"):
-                role = self.appbuilder.sm.find_role(user_details.get("base_role", "user").title())
-            else:
-                role = self.appbuilder.sm.find_role("Admin")
             user = self.appbuilder.sm.find_user(username=user_details['user_name'])
             if user:
-                # Update user role
-                user.role = role
-                self.appbuilder.sm.update_user(user=user)
+                logger.info(f"User found: {user}")
                 login_user(user, remember=False)
             else:
                 user = self.appbuilder.sm.add_user(
@@ -45,7 +39,7 @@ class DataflowAuthDBView(AuthDBView):
                     first_name=user_details.get("first_name", ""),
                     last_name=user_details.get("last_name", ""), 
                     email=user_details.get("email", ""), 
-                    role=role,
+                    role=self.appbuilder.sm.find_role(user_details.get("base_role", "user").title())
                 )
                 logger.info(f"New user created: {user}")
                 if user:

authenticator/dataflowhubauthenticator.py

@@ -1,6 +1,4 @@
-import os
-import uuid
-import re
+import os, uuid, re, hashlib, secrets
 from datetime import datetime, timedelta
 from zoneinfo import ZoneInfo
 from traitlets import Bool, Unicode
@@ -8,7 +6,8 @@ from jupyterhub.auth import Authenticator
 from oauthenticator.google import GoogleOAuthenticator
 from oauthenticator.azuread import AzureAdOAuthenticator
 from dataflow.db import get_db
-from dataflow.models import user as m_user, session as m_session, role as m_role
+from dataflow.models import user as m_user, session as m_session
+from sqlalchemy import or_
 
 class DataflowBaseAuthenticator(Authenticator):
     enable_dataflow_auth = Bool(True, config=True, help="Enable username/password authentication")
@@ -28,7 +27,7 @@ class DataflowBaseAuthenticator(Authenticator):
         return str(uuid.uuid4())
 
     def set_session_cookie(self, handler, session_id):
-        expires = datetime.now(ZoneInfo("UTC")) + timedelta(days=365)
+        expires = datetime.now(ZoneInfo("UTC")) + timedelta(days=60)
         host = handler.request.host
         domain = '.'.join(host.split('.')[-2:]) if len(host.split('.')) >= 2 else host
         handler.set_cookie(
@@ -44,19 +43,12 @@ class DataflowBaseAuthenticator(Authenticator):
         self.log.info(f"Set session cookie: dataflow_session={session_id} for host={host}")
 
     def get_or_create_session(self, user_id):
-        existing_session = (
-            self.db.query(m_session.Session)
-            .filter(m_session.Session.user_id == str(user_id))
-            .first()
-        )
-        if existing_session:
-            self.log.info(f"Reusing existing session: {existing_session.session_id}")
-            return existing_session.session_id
         session_id = self.generate_session_id()
         while self.db.query(m_session.Session).filter(
             m_session.Session.session_id == session_id
         ).first():
             session_id = self.generate_session_id()
+
         db_item = m_session.Session(user_id=user_id, session_id=session_id)
         self.db.add(db_item)
         self.db.commit()
@@ -77,42 +69,28 @@ class DataflowBaseAuthenticator(Authenticator):
 
         return super().check_blocked_users(username, authenticated)
 
-    def get_applicant_role_id(self):
-        """Get the role ID for 'Applicant' role"""
-        try:
-            applicant_role = (
-                self.db.query(m_role.Role)
-                .filter(m_role.Role.name == "Applicant")
-                .first()
-            )
-            if applicant_role:
-                return applicant_role.id
-            else:
-                self.log.warning("Applicant role not found in database")
-                return None
-        except Exception as e:
-            self.log.error(f"Error getting Applicant role: {str(e)}")
-            return None
-
     def extract_username_from_email(self, email):
         """Extract username from email by removing domain"""
         if '@' in email:
             return email.split('@')[0]
         return email
 
+    def generate_secure_password(self):
+        """Generate a secure random password hash"""
+        salt = secrets.token_hex(16)
+        random_uuid = str(uuid.uuid4())
+        hash_obj = hashlib.sha256((random_uuid + salt).encode())
+        return hash_obj.hexdigest()
+
     def create_new_user(self, email, first_name=None, last_name=None):
         """Create a new user with Applicant role"""
         try:
-            role_id = self.get_applicant_role_id()
-            if not role_id:
-                self.log.error("Cannot create user: Applicant role not found")
-                return None
-
             username = self.extract_username_from_email(email)
             username = re.sub(r'[^a-z0-9]', '', username.lower())
             if not username:
                 self.log.error("Cannot create user: Username is empty")
                 return None
+            
             existing_user = (
                 self.db.query(m_user.User)
                 .filter(m_user.User.user_name == username)
@@ -122,7 +100,7 @@ class DataflowBaseAuthenticator(Authenticator):
                 counter = 1
                 original_username = username
                 while existing_user:
-                    username = f"{original_username}_{counter}"
+                    username = f"{original_username}{counter}"
                     existing_user = (
                         self.db.query(m_user.User)
                         .filter(m_user.User.user_name == username)
@@ -130,13 +108,13 @@ class DataflowBaseAuthenticator(Authenticator):
                     )
                     counter += 1
 
+            secure_password = self.generate_secure_password()
             new_user = m_user.User(
                 user_name=username,
                 first_name=first_name or username,
                 last_name=last_name or "",
                 email=email,
-                role_id=role_id,
-                password='user@123',
+                password=secure_password,
             )
             
             self.db.add(new_user)
@@ -154,22 +132,29 @@ class DataflowBaseAuthenticator(Authenticator):
     async def authenticate_dataflow(self, handler, data):
         if not (self.enable_dataflow_auth and isinstance(data, dict) and data.get("username") and data.get("password")):
             return None
-        username = data["username"]
+        user_name_or_email = data["username"]
         password = data["password"]
-        self.log.info(f"Attempting Dataflow authentication for user: {username}")
+        self.log.info(f"Attempting Dataflow authentication for user: {user_name_or_email}")
         try:
             user = (
                 self.db.query(m_user.User)
-                .filter(m_user.User.user_name == username)
+                .filter(
+                    or_(
+                        m_user.User.email == user_name_or_email,
+                        m_user.User.user_name == user_name_or_email
+                    )
+                )
                 .first()
             )
+
             if not user or user.password != password:
-                self.log.warning(f"Dataflow authentication failed for user: {username}")
+                self.log.warning(f"Dataflow authentication failed for user: {user_name_or_email}")
                 return None
+            
             session_id = self.get_or_create_session(user.user_id)
             self.set_session_cookie(handler, session_id)
-            self.log.info(f"Dataflow authentication successful for user: {username}")
-            return {"name": username, "session_id": session_id, "auth_state": {}}
+            self.log.info(f"Dataflow authentication successful for user: {user.user_name}")
+            return {"name": user.user_name, "session_id": session_id, "auth_state": {}}
         except Exception as e:
             self.log.error(f"Dataflow authentication error: {str(e)}")
             return None

dataflow/dataflow.py

@@ -25,16 +25,13 @@ class Dataflow:
     def _parse_response_data(self, response):
         """Parse response data based on datatype field or fallback to JSON parsing."""
         data = response.json()
+        if not isinstance(data, dict):
+            raise ValueError("Internal Dataflow Error!")
         value = data.get('value', '')
-        if isinstance(data, dict) and 'datatype' in data:
-            value = data.get('value', '')
-            datatype = data.get('datatype')
-            if datatype == 'json':
-                return self._json_parse(value)
-            else:
-                return value
-        else:
+        if data.get('datatype') == 'json':
             return self._json_parse(value)
+        else:
+            return value
 
     def auth(self, session_id: str):
         """
@@ -84,12 +81,19 @@ class Dataflow:
             host_name = os.environ.get("HOSTNAME", "")
             runtime = os.environ.get("RUNTIME")
             slug = os.environ.get("SLUG")
+            org_id = os.environ.get("ORGANIZATION")
 
             dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            query_params = {
+                "key": variable_name,
+            }
 
             variable_api = None
             if runtime and slug:
                 variable_api = dataflow_config.get_config_value("auth", "variable_ui_api")
+                query_params["runtime"] = runtime
+                query_params["slug"] = slug
+                query_params["org_id"] = org_id
             elif host_name:
                 variable_api = dataflow_config.get_config_value("auth", "variable_manager_api")
             else:
@@ -98,36 +102,13 @@ class Dataflow:
             if not variable_api:
                 print("[Dataflow.variable] Variable Unreachable")
                 return None
-
-            if runtime:
-                query_params = {
-                    "key": variable_name,
-                    "runtime": runtime,
-                    "slug": slug
-                }
-                response = requests.get(variable_api, params=query_params)
-                if response.status_code == 200:
-                    response_text = response.text.strip().strip('"')
-                    return response_text
-                
-                query_params["slug"] = "global"
-                response = requests.get(variable_api, params=query_params)
-                if response.status_code == 200:
-                    response_text = response.text.strip().strip('"')
-                    return response_text
-                else: 
-                    return None
-
-            query_params = {
-                "key": variable_name,
-            }
+        
             response = requests.get(variable_api, params=query_params)
             
-            # Handle different HTTP status codes gracefully
             if response.status_code == 404:
-                return None  # Variable not found
+                return None
             elif response.status_code >= 500:
-                response.raise_for_status()  # Let server errors propagate
+                response.raise_for_status()
             elif response.status_code >= 400:
                 print(f"[Dataflow.variable] Client error {response.status_code} for variable '{variable_name}'")
                 return None
@@ -158,32 +139,30 @@ class Dataflow:
             host_name = os.environ.get("HOSTNAME", "")
             runtime = os.environ.get("RUNTIME")
             slug = os.environ.get("SLUG")
+            org_id = os.environ.get("ORGANIZATION")
 
             dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
-            if runtime:
-                secret_api = dataflow_config.get_config_value("auth", "secret_ui_api")
-            else:
-                secret_api = dataflow_config.get_config_value("auth", "secret_manager_api")
-            if not secret_api:
-                print("[Dataflow.secret] Secret API Unreachable")
-                return None
-
             query_params = {
                 "key": secret_name
             }
 
             if runtime:
+                secret_api = dataflow_config.get_config_value("auth", "secret_ui_api")
                 query_params["runtime"] = runtime
-            if slug:
                 query_params["slug"] = slug
+                query_params["org_id"] = org_id
+            else:
+                secret_api = dataflow_config.get_config_value("auth", "secret_manager_api")
+            if not secret_api:
+                print("[Dataflow.secret] Secret API Unreachable")
+                return None
 
             response = requests.get(secret_api, params=query_params)
-            
-            # Handle different HTTP status codes gracefully
+
             if response.status_code == 404:
-                return None  # Secret not found
+                return None
             elif response.status_code >= 500:
-                response.raise_for_status()  # Let server errors propagate
+                response.raise_for_status()
             elif response.status_code >= 400:
                 print(f"[Dataflow.secret] Client error {response.status_code} for secret '{secret_name}'")
                 return None
@@ -214,31 +193,29 @@ class Dataflow:
             host_name = os.environ["HOSTNAME"]
             runtime = os.environ.get("RUNTIME")
             slug = os.environ.get("SLUG")
-            
+            org_id = os.environ.get("ORGANIZATION")
+
             dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
-            if runtime:
-                connection_api = dataflow_config.get_config_value("auth", "connection_ui_api")
-            elif host_name:
-                connection_api = dataflow_config.get_config_value("auth", "connection_manager_api")
-            else:
-                raise Exception("Cannot run dataflow methods here! HOSTNAME or RUNTIME env variable not set.")
-            
             query_params = {
                 "conn_id": conn_id
             }
 
             if runtime:
                 query_params["runtime"] = runtime
-            if slug:
+                query_params["org_id"] = org_id
                 query_params["slug"] = slug
+                connection_api = dataflow_config.get_config_value("auth", "connection_ui_api")
+            elif host_name:
+                connection_api = dataflow_config.get_config_value("auth", "connection_manager_api")
+            else:
+                raise Exception("Cannot run dataflow methods here! HOSTNAME or RUNTIME env variable not set.")
 
             response = requests.get(connection_api, params=query_params)
             
-            # Handle different HTTP status codes gracefully
             if response.status_code == 404:
                 raise RuntimeError(f"[Dataflow.connection] Connection '{conn_id}' not found!")
             elif response.status_code >= 500:
-                response.raise_for_status()  # Let server errors propagate
+                response.raise_for_status()
             elif response.status_code >= 400:
                 raise RuntimeError(f"[Dataflow.connection] Client error {response.status_code} for connection '{conn_id}'")
             elif response.status_code != 200:
@@ -306,20 +283,20 @@ class Dataflow:
             host_name = os.environ.get("HOSTNAME", "")
             runtime = os.environ.get("RUNTIME")
             slug = os.environ.get("SLUG")
+            org_id = os.environ.get("ORGANIZATION")
 
             dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
-            if runtime and slug:
-                variableorsecret_api = dataflow_config.get_config_value("auth", "variableorsecret_ui_api")
-                query_params = {
-                    "key": key,
-                    "runtime": runtime,
-                    "slug": slug
+            query_params = {
+                    "key": key
                 }
+            
+            if runtime:
+                variableorsecret_api = dataflow_config.get_config_value("auth", "variableorsecret_ui_api")
+                query_params["runtime"] = runtime
+                query_params["slug"] = slug
+                query_params["org_id"] = org_id
             elif host_name:
                 variableorsecret_api = dataflow_config.get_config_value("auth", "variableorsecret_manager_api")
-                query_params = {
-                    "key": key
-                }
             else:
                 raise Exception("Cannot run dataflow methods here!") 
 
@@ -329,9 +306,8 @@ class Dataflow:
             
             response = requests.get(variableorsecret_api, params=query_params)
             
-            # Handle different HTTP status codes gracefully
             if response.status_code == 404:
-                return None  # Variable/secret not found
+                return None
             elif response.status_code >= 500:
                 response.raise_for_status()  # Let server errors propagate
             elif response.status_code >= 400:

dataflow/environment.py

@@ -6,12 +6,14 @@ from .configuration import ConfigurationManager
 from .utils.logger import CustomLogger
 
 class EnvironmentManager:
-    def __init__(self):
+    def __init__(self, org_id: int = None):
         """Initialize the EnvironmentManager"""
         self.config = ConfigurationManager('/dataflow/app/config/dataflow.cfg')
-        self.env_base_path = self.config.get_config_value('paths', 'env_path')
-        self.env_logs_path = self.config.get_config_value('paths', 'env_logs_path')
-        self.env_version_path = self.config.get_config_value('paths', 'env_versions_path')
+        self.org_id = org_id
+        self.env_sub_path = f"{self.org_id}" if self.org_id is not None else "dataflow"
+        self.env_base_path = os.path.join(self.config.get_config_value('paths', 'env_path'), self.env_sub_path, 'python_envs')
+        self.env_logs_path = os.path.join(self.config.get_config_value('paths', 'env_logs_path'), self.env_sub_path, 'logs')
+        self.env_version_path = os.path.join(self.config.get_config_value('paths', 'env_versions_path'), self.env_sub_path, 'versions')
         self.local_env_logs_path = self.config.get_config_value('paths', 'local_env_logs_path')
         os.makedirs(self.env_version_path, exist_ok=True)
         self.logger = CustomLogger().get_logger(__name__)
@@ -334,10 +336,9 @@ class EnvironmentManager:
         """
         versioned_name = f"{env_name}_v{env_version}"
         log_file_name = f"envlog_{versioned_name}.log"
-        log_file_dir = self.config.get_config_value('paths', 'env_logs_path')
-        os.makedirs(log_file_dir, exist_ok=True)
-        log_file_location = os.path.join(log_file_dir, log_file_name)
-        
+        os.makedirs(self.env_logs_path, exist_ok=True)
+        log_file_location = os.path.join(self.env_logs_path, log_file_name)
+
         # Clear log file if it exists
         if os.path.exists(log_file_location):
             open(log_file_location, "w").close()
@@ -394,8 +395,14 @@ class EnvironmentManager:
         Returns:
             JobLogs: The created or updated job entry
         """
-        job = db.query(JobLogs).filter(JobLogs.log_file_name == log_file_name).first()
-
+        job = (
+            db.query(JobLogs)
+            .filter(
+                JobLogs.log_file_name == log_file_name,
+                JobLogs.org_id == self.org_id if self.org_id is not None else JobLogs.org_id.is_(None)
+            )
+            .first()
+        )
         if job:
             if job.status == "success":
                 self.logger.error(f"Job with log_file_name '{log_file_name}' already completed successfully.")
@@ -409,6 +416,7 @@ class EnvironmentManager:
                 log_file_name=log_file_name,
                 log_file_location=log_file_location,
                 created_by=user_name,
+                org_id=self.org_id if self.org_id else None,
                 status="in_progress"
             )
             db.add(job)
@@ -464,7 +472,8 @@ class EnvironmentManager:
             env_status = "Draft" if status == "success" else "Failed"
 
             db.query(Environment).filter(
-                Environment.short_name == env_short_name
+                Environment.short_name == env_short_name,
+                Environment.org_id == self.org_id if self.org_id is not None else Environment.org_id.is_(None)
             ).update({"version": version, "pip_libraries": pip_libraries, "conda_libraries": conda_libraries, "status": env_status})
             db.commit()
 

dataflow/models/__init__.py

@@ -1,9 +1,9 @@
 # init for loading models in the application
 
 from .role import Role
-from .user import User
+from .user import User, UserOnboarding, OnboardingStatus
 from .team import Team
-from .environment import (Environment, LocalEnvironment, ArchivedEnvironment, JobLogs)
+from .environment import (Environment, LocalEnvironment, ArchivedEnvironment, JobLogs, PipSource)
 from .project_details import ProjectDetails
 from .recent_projects import RecentProjects
 from .pinned_projects import PinnedProject
@@ -22,4 +22,6 @@ from .recent_project_studio import RecentProjectStudio
 from .connection import Connection
 from .git_ssh import GitSSH
 from .pod_activity import PodActivity
-from .pod_session_history import PodSessionHistory
+from .pod_session_history import PodSessionHistory
+from .organization import Organization, OrganizationOnboarding
+from .org_associations import OrganizationServer, OrganizationUser, OrganizationAppType

dataflow/models/app_types.py

@@ -1,10 +1,15 @@
 from sqlalchemy import Column, Integer, String, Boolean
+from sqlalchemy.orm import relationship
 from dataflow.db import Base
 
 class AppType(Base):
-    __tablename__ = "RUNTIME_APP_TYPE"
+    __tablename__ = "APP_TYPE"
     
     id = Column(Integer, primary_key=True, autoincrement=True, unique=True)
-    name = Column(String, unique=True, nullable=True)
+    name = Column(String, unique=True, nullable=False)
     display_name = Column(String, nullable=False)
-    code_based = Column(Boolean, nullable=False)
+    code_based = Column(Boolean, nullable=False)
+    studio = Column(Boolean, nullable=False, default=False, server_default='false')
+    runtime = Column(Boolean, nullable=False, default=False, server_default='false')
+
+    organizations = relationship("Organization", secondary="ORGANIZATION_APP_TYPE", back_populates="apps")

dataflow/models/connection.py

@@ -1,4 +1,4 @@
-from sqlalchemy import Column, String, Integer, Boolean, DateTime, UniqueConstraint
+from sqlalchemy import Column, String, Integer, Boolean, DateTime, UniqueConstraint, ForeignKey
 from sqlalchemy.sql import func
 from dataflow.db import Base
 
@@ -10,16 +10,17 @@ class Connection(Base):
 
     id = Column(Integer, primary_key=True, index=True)
     conn_id = Column(String, index=True, nullable=False)
+    org_id = Column(Integer, ForeignKey("ORGANIZATION.id"), index=True, nullable=False)
     description = Column(String, nullable=True)
     conn_type = Column(String, nullable=False)
     runtime = Column(String, nullable=True)
     slug = Column(String, nullable=True)
-    status = Column(Boolean, default=False)
+    status = Column(Boolean, default=False, server_default='false')
     created_by = Column(String, nullable=True)
     created_at = Column(DateTime(timezone=True), server_default=func.now())
     updated_at = Column(DateTime(timezone=True), onupdate=func.now())
-    is_active = Column(Boolean, default=True)
+    is_active = Column(Boolean, default=True, server_default='true')
     
     __table_args__ = (
-        UniqueConstraint('conn_id', 'runtime', 'slug', 'is_active', 'created_by', name='uq_active_conn_with_runtime_slug'),
+        UniqueConstraint('conn_id', 'org_id', 'runtime', 'slug', 'is_active', 'created_by', name='uq_active_conn_with_runtime_slug'),
     )

dataflow/models/dataflow_zone.py

@@ -8,10 +8,9 @@ class DataflowZone(Base):
     id = Column(Integer, primary_key=True, autoincrement=True)
     slug = Column(String, unique=True, nullable=False)
     display_name = Column(String, nullable=False)
-    is_runtime = Column(Boolean, default=False)
+    is_runtime = Column(Boolean, default=False, server_default='false')
     subdomain = Column(String)
-    spark_enabled = Column(Boolean, default=False)
-    display_order = Column(Integer, default=0)
+    display_order = Column(Integer, default=0, server_default='0')
 
     role_zone_assocs = relationship("RoleZone", back_populates="zone")