dataflow-core 2.1.6__tar.gz → 2.1.8__tar.gz

{dataflow_core-2.1.6 → dataflow_core-2.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dataflow-core
-Version: 2.1.6
+Version: 2.1.8
 Summary: Dataflow core package
 Author: Dataflow
 Author-email: 
@@ -9,6 +9,8 @@ Requires-Dist: boto3
 Requires-Dist: psycopg2-binary
 Requires-Dist: pymysql
 Requires-Dist: requests
+Requires-Dist: azure-identity
+Requires-Dist: azure-keyvault-secrets
 Dynamic: author
 Dynamic: requires-dist
 Dynamic: summary

{dataflow_core-2.1.6 → dataflow_core-2.1.8}/authenticator/dataflowhubauthenticator.py

@@ -109,7 +109,7 @@ class DataflowBaseAuthenticator(Authenticator):
                 return None
 
             username = self.extract_username_from_email(email)
-            username = re.sub(r'[^A-Za-z0-9]', '', username)
+            username = re.sub(r'[^a-z0-9]', '', username.lower())
             if not username:
                 self.log.error("Cannot create user: Username is empty")
                 return None
@@ -136,7 +136,6 @@ class DataflowBaseAuthenticator(Authenticator):
                 last_name=last_name or "",
                 email=email,
                 role_id=role_id,
-                active='Y',
                 password='user@123',
             )
             
@@ -248,7 +247,11 @@ class DataflowAzureAuthenticator(DataflowBaseAuthenticator, AzureAdOAuthenticato
     azure_client_secret = Unicode(config=True, help="Azure AD OAuth client secret")
     azure_tenant_id = Unicode(config=True, help="Azure AD tenant ID")
     azure_scope = Unicode("openid profile email", config=True, help="Azure AD OAuth scopes")
-
+    dataflow_oauth_type = Unicode(
+        default_value="google",
+        config=True,
+        help="The OAuth provider type for DataflowHub (e.g., github, google)"
+    )
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self.client_id = self.azure_client_id
@@ -270,48 +273,47 @@ class DataflowAzureAuthenticator(DataflowBaseAuthenticator, AzureAdOAuthenticato
             if not user:
                 self.log.warning("Azure AD OAuth authentication failed: No user data returned")
                 return None
-            
-            email = user.get("email") or user.get("preferred_username")
+
+            auth_state = user.get("auth_state", {})
+            user_info = auth_state.get("user", {}) if auth_state else {}
+            email = user_info.get("upn")
             if not email:
-                self.log.warning("Azure AD OAuth authentication failed: No email in user data")
+                self.log.warning("Azure AD OAuth authentication failed: No upn in user data")
                 return None
-            
+
             db_user = (
                 self.db.query(m_user.User)
                 .filter(m_user.User.email == email)
                 .first()
             )
-            
+
             if not db_user:
                 self.log.info(f"User with email {email} not found in Dataflow database, creating new user")
-                # Extract additional info from user data if available
-                auth_state = user.get("auth_state", {})
-                user_info = auth_state.get("user", {}) if auth_state else {}
                 
-                first_name = user_info.get("given_name") or user.get("given_name")
-                last_name = user_info.get("family_name") or user.get("family_name")
+                first_name = user_info.get("name") or user.get("name")
                 
-                db_user = self.create_new_user(email, first_name, last_name)
+                db_user = self.create_new_user(email, first_name, last_name=None)
                 if not db_user:
                     self.log.error(f"Failed to create new user for email: {email}")
                     return None
-            
+
             username = db_user.user_name
             session_id = self.get_or_create_session(db_user.user_id)
             self.set_session_cookie(handler, session_id)
             self.log.info(f"Azure AD OAuth completed for user: {username}, session_id={session_id}")
             return {
-                "name": username,
+                "name": db_user.first_name,
                 "session_id": session_id,
                 "auth_state": user.get("auth_state", {})
             }
+
         except Exception as e:
             self.log.error(f"Azure AD OAuth authentication error: {str(e)}", exc_info=True)
             return None
         finally:
             self.db.close()
 
-auth_type = os.environ.get("DATAFLOW_AUTH_TYPE", "google")
+auth_type = os.environ.get("DATAFLOW_OAUTH_TYPE", "google")
 
 if auth_type == "google":
     BaseAuthenticator = DataflowGoogleAuthenticator

dataflow_core-2.1.8/dataflow/dataflow.py

@@ -0,0 +1,323 @@
+import os, requests
+from .database_manager import DatabaseManager
+import json
+from .configuration import ConfigurationManager
+
+
+class Dataflow:
+    """
+    Dataflow class to interact with Dataflow services.
+    """
+    def auth(self, session_id: str):
+        """
+        Retrieve and return user information using their session ID.
+        
+        Args:
+            session_id (str): User's session ID from cookies
+            
+        Returns:
+            dict: User information including username, name, email, and role
+        """
+        try:
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+            auth_api = dataflow_config.get_config_value('auth', 'ui_auth_api')
+            response = requests.get(
+                auth_api,
+                cookies={"dataflow_session": session_id, "jupyterhub-hub-login": ""}
+            )
+            
+            if response.status_code != 200:
+                return response.json()
+            
+            user_data = response.json()
+            user_dict = {
+                "user_name": user_data["user_name"], 
+                "first_name": user_data["first_name"],
+                "last_name": user_data["last_name"] if user_data.get("last_name") else "",
+                "email": user_data["email"],
+                "role": user_data["base_role"]
+            }
+            return user_dict
+                  
+        except Exception as e:
+            return e
+    
+    def variable(self, variable_name: str):
+        """
+        Retrieve a Dataflow variable.
+        
+        Args:
+            variable_name (str): Name of the variable to retrieve
+            
+        Returns:
+            str or None: Variable value if found, None otherwise
+        """
+        try:
+            host_name = os.environ.get("HOSTNAME", "")
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+
+            dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+
+            variable_api = None
+            if runtime and slug:
+                variable_api = dataflow_config.get_config_value("auth", "variable_ui_api")
+            elif host_name:
+                variable_api = dataflow_config.get_config_value("auth", "variable_manager_api")
+            else:
+                raise Exception("Cannot run dataflow methods here!") 
+            
+            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
+            elif response.status_code >= 500:
+                response.raise_for_status()  # Let server errors propagate
+            elif response.status_code >= 400:
+                print(f"[Dataflow.variable] Client error {response.status_code} for variable '{variable_name}'")
+                return None
+            elif response.status_code != 200:
+                print(f"[Dataflow.variable] Unexpected status {response.status_code} for variable '{variable_name}'")
+                return None
+
+            return response.text.strip().strip('"')
+        
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"[Dataflow.variable] Failed to fetch variable '{variable_name}'") from e
+        
+        except Exception as e:
+            print(f"[Dataflow.variable] Exception occurred: {e}")
+            return None
+        
+    def secret(self, secret_name: str):
+        """
+        Retrieve a Dataflow secret value.
+        
+        Args:
+            secret_name (str): Name of the secret to retrieve
+            
+        Returns:
+            str or None: Secret value if found, None otherwise
+        """
+        try:
+            host_name = os.environ.get("HOSTNAME", "")
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+
+            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:
+                query_params["runtime"] = runtime
+            if slug:
+                query_params["slug"] = slug
+
+            response = requests.get(secret_api, params=query_params)
+            
+            # Handle different HTTP status codes gracefully
+            if response.status_code == 404:
+                return None  # Secret not found
+            elif response.status_code >= 500:
+                response.raise_for_status()  # Let server errors propagate
+            elif response.status_code >= 400:
+                print(f"[Dataflow.secret] Client error {response.status_code} for secret '{secret_name}'")
+                return None
+            elif response.status_code != 200:
+                print(f"[Dataflow.secret] Unexpected status {response.status_code} for secret '{secret_name}'")
+                return None
+
+            return response.text.strip().strip('"')
+        
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"[Dataflow.secret] Failed to fetch secret '{secret_name}'") from e
+        except Exception as e:
+            print(f"[Dataflow.secret] Exception occurred: {e}")
+            return None
+
+    def connection(self, conn_id: str, mode="session"):
+        """
+        Connects with a Dataflow connection.
+        
+        Args:
+            conn_id (str): Connection identifier
+            mode (str): Return type - "session" (default) or "engine" or "url"
+            
+        Returns:
+            Session or Engine: SQLAlchemy session or engine based on mode
+        """
+        try:
+            host_name = os.environ["HOSTNAME"]
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+            
+            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["slug"] = slug
+
+            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
+            elif response.status_code >= 400:
+                raise RuntimeError(f"[Dataflow.connection] Client error {response.status_code} for connection '{conn_id}'")
+            elif response.status_code != 200:
+                raise RuntimeError(f"[Dataflow.connection] Unexpected status {response.status_code} for connection '{conn_id}'")
+                
+            connection_details = response.json()
+
+            if not connection_details:
+                raise RuntimeError(f"[Dataflow.connection] Connection '{conn_id}' not found!")
+            
+            if mode == "dict":
+                with open('/home/jovyan/log.txt', 'w') as log_file:
+                    log_file.write(f"Connection details for {conn_id}: {connection_details}\n")
+                print(f"connection_details: {connection_details}")
+                return dict(connection_details)
+
+            conn_type = connection_details['conn_type'].lower()
+            username = connection_details['login']
+            password = connection_details.get('password', '')
+            host = connection_details['host']
+            port = connection_details['port']
+            database = connection_details.get('schemas', '')
+
+            user_info = f"{username}:{password}@" if password else f"{username}@"
+            db_info = f"/{database}" if database else ""
+
+            connection_string = f"{conn_type}://{user_info}{host}:{port}{db_info}"
+
+            extra = connection_details.get('extra', '')            
+            if extra:
+                try:
+                    extra_params = json.loads(extra)
+                    if extra_params:
+                        extra_query = "&".join(f"{key}={value}" for key, value in extra_params.items())
+                        connection_string += f"?{extra_query}"
+                except json.JSONDecodeError:
+                    # If 'extra' is not valid JSON, skip adding extra parameters
+                    pass
+
+            if mode == "url":
+                return connection_string
+    
+            connection_instance = DatabaseManager(connection_string)
+            if mode == "engine":
+                return connection_instance.get_engine()
+            elif mode == "session":
+                return next(connection_instance.get_session())
+            else:
+                raise ValueError(f"Unsupported mode: {mode}. Use 'session', 'engine', 'url'.")
+        
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"[Dataflow.connection] Failed to fetch connection '{conn_id}'") from e
+
+        except Exception as e:
+            raise RuntimeError(f"[Dataflow.connection] Error connecting to '{conn_id}': {str(e)}") from e
+        
+    def variable_or_secret(self, key: str):
+        """
+        Retrieve a variable or secret by key.
+        
+        Args:
+            key (str): Key of the variable or secret
+            
+        Returns:
+            str or None: Value if found, None otherwise
+        """
+        try:
+            host_name = os.environ.get("HOSTNAME", "")
+            runtime = os.environ.get("RUNTIME")
+            slug = os.environ.get("SLUG")
+
+            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
+                }
+            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!") 
+
+            if not variableorsecret_api:
+                print("[Dataflow.variable_or_secret] Variable/Secret Unreachable")
+                return None
+            
+            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
+            elif response.status_code >= 500:
+                response.raise_for_status()  # Let server errors propagate
+            elif response.status_code >= 400:
+                print(f"[Dataflow.variable_or_secret] Client error {response.status_code} for key '{key}'")
+                return None
+            elif response.status_code != 200:
+                print(f"[Dataflow.variable_or_secret] Unexpected status {response.status_code} for key '{key}'")
+                return None
+
+            response_text = response.text.strip().strip('"')
+            return response_text
+            
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"[Dataflow.variable_or_secret] Failed to fetch '{key}'") from e

{dataflow_core-2.1.6 → dataflow_core-2.1.8}/dataflow/models/__init__.py

@@ -12,7 +12,8 @@ from .blacklist_library import BlacklistedLibrary
 from .environment_status import EnvironmentStatus
 from .session import Session
 from .server_config import ServerConfig, CustomServerConfig
-from .runtime import RuntimeZone
+from .dataflow_zone import DataflowZone
+from .role_zone import RoleZone
 from .environment_status import EnvironmentStatus
 from .user_team import UserTeam
 from .role_server import RoleServer

dataflow_core-2.1.8/dataflow/models/dataflow_zone.py

@@ -0,0 +1,19 @@
+from sqlalchemy import Column, Integer, String, Boolean
+from sqlalchemy.orm import relationship
+from dataflow.db import Base
+
+class DataflowZone(Base):
+    __tablename__ = "DATAFLOW_ZONE"
+
+    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)
+    subdomain = Column(String)
+    spark_enabled = Column(Boolean, default=False)
+    display_order = Column(Integer, default=0)
+
+    role_zone_assocs = relationship("RoleZone", back_populates="zone")
+
+    def __repr__(self):
+        return f"<DataflowZone(id={self.id}, slug='{self.slug}', display_name='{self.display_name}', display_order={self.display_order})>"

{dataflow_core-2.1.6 → dataflow_core-2.1.8}/dataflow/models/role.py

@@ -2,6 +2,12 @@
 from sqlalchemy import Column, Integer, String, Enum
 from sqlalchemy.orm import relationship
 from dataflow.db import Base
+import enum
+
+class BaseRoleField(enum.Enum):
+    admin = "admin"
+    user = "user"
+    applicant = "applicant"
 
 class Role(Base):
     """
@@ -13,7 +19,11 @@ class Role(Base):
     id = Column(Integer, primary_key=True, index=True, autoincrement=True, nullable=False)
     name = Column(String, unique=True, nullable=False)
     description = Column(String, nullable=True)
-    base_role = Column(Enum('admin', 'user', 'applicant', name='base_role_field'), default='user', nullable=False)
+    base_role = Column(Enum(BaseRoleField), nullable=False, default=BaseRoleField.user)
 
     users = relationship("User", back_populates="role_details", cascade="all, delete-orphan")
-    role_server_assocs = relationship("RoleServer", back_populates="role")
+    role_server_assocs = relationship("RoleServer", back_populates="role")
+    role_zone_assocs = relationship("RoleZone", back_populates="role")
+
+    def __repr__(self):
+        return f"<Role(id={self.id}, name='{self.name}', base_role='{self.base_role}')>"

dataflow_core-2.1.8/dataflow/models/role_zone.py

@@ -0,0 +1,17 @@
+from typing import Dict, List, Optional
+from sqlalchemy import Column, Integer, ForeignKey, UniqueConstraint, Boolean
+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)
+
+    role = relationship("Role", back_populates="role_zone_assocs")
+    zone = relationship("DataflowZone", back_populates="role_zone_assocs")
+
+    def __repr__(self):
+        return f"<RoleZone(role_id={self.role_id}, zone_id={self.zone_id}, is_default={self.is_default})>"

{dataflow_core-2.1.6 → dataflow_core-2.1.8}/dataflow/models/user.py

@@ -1,5 +1,5 @@
 """models.py"""
-from sqlalchemy import Column, Integer, String, Boolean, LargeBinary, Enum, ForeignKey
+from sqlalchemy import Column, Integer, String, Boolean, LargeBinary, ForeignKey
 from sqlalchemy.orm import relationship
 from dataflow.db import Base
 
@@ -18,7 +18,7 @@ class User(Base):
     role_id = Column(Integer, ForeignKey('ROLE.id'), nullable=False)
     image = Column(LargeBinary)
     image_url = Column(String, nullable=True)
-    active = Column(Enum('N', 'Y', name='active_field'), nullable=False, server_default=str("N"))
+    active = Column(Boolean, nullable=False, default=True)
     password = Column(String, nullable=False)
     active_env = Column(String)
     active_env_type = Column(String, nullable=True)

dataflow_core-2.1.8/dataflow/schemas/connection.py

@@ -0,0 +1,84 @@
+"""schemas/connection.py"""
+
+from pydantic import BaseModel, field_validator
+from typing import Optional
+from datetime import datetime
+from enum import Enum
+
+
+class ConnectionType(str, Enum):
+    """Enum for supported connection types."""
+    POSTGRESQL = "PostgreSQL"
+    MYSQL = "MySQL"
+
+
+class ConnectionBase(BaseModel):
+    """Base connection model with common fields."""
+    conn_id: str
+    conn_type: ConnectionType
+    description: Optional[str] = None
+    host: str
+    schemas: Optional[str] = None
+    password: str
+    login: str
+    port: int
+    extra: Optional[str] = None
+
+    @field_validator("conn_id")
+    def validate_conn_id(cls, v) -> str:
+        import re
+        if not isinstance(v, str):
+            raise ValueError("Connection ID must be a string.")
+        if len(v) > 20:
+            raise ValueError("Connection ID must be at most 20 characters long.")
+        if not re.fullmatch(r"[A-Za-z0-9-]+", v):
+            raise ValueError(
+                "Connection ID can only contain letters, numbers, and hyphens (-)!"
+            )
+        return v
+
+    @field_validator("conn_type")
+    def validate_conn_type(cls, v) -> ConnectionType:
+        if isinstance(v, str):
+            try:
+                return ConnectionType(v)
+            except ValueError:
+                raise ValueError(f'conn_type must be one of {[e.value for e in ConnectionType]}')
+        return v
+
+
+class ConnectionSave(ConnectionBase):
+    """Model for creating a new connection."""
+    pass
+
+
+class ConnectionUpdate(BaseModel):
+    """Model for updating an existing connection."""
+    conn_type: Optional[ConnectionType] = None
+    description: Optional[str] = None
+    host: Optional[str] = None
+    schemas: Optional[str] = None
+    login: Optional[str] = None
+    password: Optional[str] = None
+    port: Optional[int] = None
+    extra: Optional[str] = None
+
+    @field_validator("conn_type")
+    def validate_conn_type(cls, v) -> Optional[ConnectionType]:
+        if v is None:
+            return v
+        if isinstance(v, str):
+            # Convert string to enum if needed
+            try:
+                return ConnectionType(v)
+            except ValueError:
+                raise ValueError(f'conn_type must be one of {[e.value for e in ConnectionType]}')
+        return v
+
+
+class ConnectionRead(ConnectionBase):
+    """Model for reading/displaying connection data."""
+    pass
+
+    class Config:
+        from_attributes = True

dataflow_core-2.1.8/dataflow/schemas/git_ssh.py

@@ -0,0 +1,50 @@
+"""schemas/git_ssh.py"""
+
+from pydantic import BaseModel, field_validator
+from typing import Optional, Literal
+from datetime import datetime
+
+
+class SSHBase(BaseModel):
+    """Base SSH key model with common fields."""
+    key_name: str
+    description: Optional[str] = None
+
+    @field_validator("key_name")
+    def validate_key_name(cls, v) -> str:
+        import re
+        if not isinstance(v, str):
+            raise ValueError("SSH key name must be a string.")
+        if len(v) > 20:
+            raise ValueError("SSH key name must be at most 20 characters long.")
+        if not re.fullmatch(r"[A-Za-z0-9-]+", v):
+            raise ValueError(
+                "SSH key name can only contain letters, numbers, and hyphens (-)!"
+            )
+        return v
+
+
+class SSHSave(SSHBase):
+    """Model for creating a new SSH key."""
+    public_key: str
+    private_key: str
+
+
+class SSHUpdate(BaseModel):
+    """Model for updating an existing SSH key."""
+    description: Optional[str] = None
+    public_key: Optional[str] = None
+    private_key: Optional[str] = None
+
+
+class SSHRead(SSHBase):
+    """Model for reading/displaying SSH key data."""
+    public_key: str
+    private_key: str
+    created_date: Optional[datetime] = None
+
+    class Config:
+        from_attributes = True
+
+    class Config:
+        from_attributes = True