dataflow-core 2.1.6__py3-none-any.whl → 2.1.8__py3-none-any.whl

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/dataflow.py

@@ -1,14 +1,13 @@
 import os, requests
 from .database_manager import DatabaseManager
-from .utils.aws_secrets_manager import SecretsManagerClient
 import json
 from .configuration import ConfigurationManager
 
 
 class Dataflow:
-    def __init__(self):
-        self.secrets_manager = SecretsManagerClient()
-
+    """
+    Dataflow class to interact with Dataflow services.
+    """
     def auth(self, session_id: str):
         """
         Retrieve and return user information using their session ID.
@@ -55,19 +54,26 @@ class Dataflow:
         """
         try:
             host_name = os.environ.get("HOSTNAME", "")
-            user_name = host_name.replace("jupyter-", "") if host_name.startswith("jupyter-") else host_name
             runtime = os.environ.get("RUNTIME")
             slug = os.environ.get("SLUG")
 
             dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
-            variable_api = dataflow_config.get_config_value("auth", "db_get_variables")
+
+            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 = {
-                    "variable_key": variable_name,
+                    "key": variable_name,
                     "runtime": runtime,
                     "slug": slug
                 }
@@ -85,17 +91,27 @@ class Dataflow:
                     return None
 
             query_params = {
-                "variable_key": variable_name,
-                "runtime": None,
-                "slug": None,
-                "created_by": user_name
+                "key": variable_name,
             }
             response = requests.get(variable_api, params=query_params)
-            if response.status_code == 200:
-                response_text = response.text.strip().strip('"')
-                return response_text
-            else:
+            
+            # 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
@@ -112,19 +128,20 @@ class Dataflow:
         """
         try:
             host_name = os.environ.get("HOSTNAME", "")
-            user_name = host_name.replace("jupyter-", "") if host_name.startswith("jupyter-") else host_name
             runtime = os.environ.get("RUNTIME")
             slug = os.environ.get("SLUG")
 
             dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
-            secret_api = dataflow_config.get_config_value("auth", "db_get_secrets")
+            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 = {
-                "secret_key": secret_name,
-                "created_by": user_name
+                "key": secret_name
             }
 
             if runtime:
@@ -134,11 +151,22 @@ class Dataflow:
 
             response = requests.get(secret_api, params=query_params)
             
-            if response.status_code == 200:
-                response_text = response.text.strip().strip('"')
-                return response_text
-            else:
+            # 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
@@ -156,26 +184,62 @@ class Dataflow:
         """
         try:
             host_name = os.environ["HOSTNAME"]
-            user_name=host_name.replace("jupyter-","")
             runtime = os.environ.get("RUNTIME")
             slug = os.environ.get("SLUG")
             
-            vault_path = "connections"
-            secret = self.secrets_manager.get_secret_by_key(vault_path, user_name, conn_id, runtime, 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 = secret['conn_type'].lower()
-            username = secret['login']
-            password = secret.get('password', '')
-            host = secret['host']
-            port = secret['port']
-            database = secret.get('schemas', '')
+            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 = secret.get('extra', '')            
+            extra = connection_details.get('extra', '')            
             if extra:
                 try:
                     extra_params = json.loads(extra)
@@ -194,6 +258,66 @@ class Dataflow:
                 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:
-            return None
+            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/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/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/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/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/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/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/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

dataflow/schemas/secret.py

@@ -0,0 +1,44 @@
+"""schemas/secret.py"""
+
+from pydantic import BaseModel, field_validator
+from typing import Optional
+from datetime import datetime
+
+
+class SecretBase(BaseModel):
+    """Base secret model with common fields."""
+    key: str
+    value: str
+    description: Optional[str] = None
+
+    @field_validator("key")
+    def validate_key(cls, v) -> str:
+        import re
+        if not isinstance(v, str):
+            raise ValueError("Secret key must be a string.")
+        if len(v) > 20:
+            raise ValueError("Secret key must be at most 20 characters long.")
+        if not re.fullmatch(r"[A-Za-z0-9-]+", v):
+            raise ValueError(
+                "Secret key can only contain letters, numbers, and hyphens (-)!"
+            )
+        return v
+
+
+class SecretSave(SecretBase):
+    """Model for creating a new secret."""
+    pass
+
+
+class SecretUpdate(BaseModel):
+    """Model for updating an existing secret."""
+    value: Optional[str] = None
+    description: Optional[str] = None
+
+
+class SecretRead(SecretBase):
+    """Model for reading/displaying secret data."""
+    created_date: Optional[datetime] = None
+
+    class Config:
+        from_attributes = True

dataflow/secrets_manager/__init__.py

@@ -0,0 +1,13 @@
+# secrets_manager/__init__.py
+
+from .factory import get_secret_manager
+from .service import SecretsService
+
+# 1. Call the factory to get the configured low-level secret manager
+#    (e.g., an instance of AWSSecretsManager or AzureKeyVault).
+#    This happens only once when the package is first imported.
+secret_manager_instance = get_secret_manager()
+
+# 2. Create the single, high-level service instance that the rest of
+#    your application will use. It wraps the low-level instance.
+secrets_service = SecretsService(secret_manager=secret_manager_instance)