dataflow-core 2.1.7__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
@@ -247,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
@@ -269,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/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)

dataflow/secrets_manager/factory.py

@@ -0,0 +1,59 @@
+# secrets_manager/factory.py
+import os
+from .interface import SecretManager
+from .providers.aws_manager import AWSSecretsManager
+from .providers.azure_manager import AzureKeyVault
+from ..configuration import ConfigurationManager
+
+# A custom exception for clear error messages
+class SecretProviderError(Exception):
+    pass
+
+def get_secret_manager() -> SecretManager:
+    """
+    Factory function to get the configured secret manager instance.
+
+    Reads the cloud provider configuration from dataflow_auth.cfg
+    to determine which cloud provider's secret manager to instantiate.
+    """
+    try:
+        # dataflow_config = None
+        # if os.getenv('HOSTNAME'):
+        #     dataflow_config = ConfigurationManager('/dataflow/app/auth_config/dataflow_auth.cfg')
+        # else:
+        dataflow_config = ConfigurationManager('/dataflow/app/config/dataflow.cfg')
+    except Exception as e:
+        raise SecretProviderError(
+            f"Failed to read cloud provider configuration: {str(e)}. "
+            "Please check that the configuration file exists and contains the 'cloud' section."
+        )
+
+    provider = dataflow_config.get_config_value('cloudProvider', 'cloud')
+    if not provider:
+        raise SecretProviderError(
+            "The cloud provider is not configured in config file. "
+            "Please set the 'cloud' value in the 'cloud' section to 'aws' or 'azure'."
+        )
+
+    provider = provider.lower()
+    print(f"Initializing secret manager for provider: {provider}")
+
+    if provider == "aws":
+        return AWSSecretsManager()
+
+    elif provider == "azure":
+        vault_url = dataflow_config.get_config_value('cloudProvider', 'key_vault')
+        if not vault_url:
+            raise SecretProviderError(
+                "AZURE_VAULT_URL must be set when using the Azure provider."
+            )
+        return AzureKeyVault(vault_url=vault_url)
+
+    # You can easily add more providers here in the future
+    # elif provider == "gcp":
+    #     return GCPSecretManager()
+
+    else:
+        raise SecretProviderError(
+            f"Unsupported secret provider: '{provider}'. Supported providers are: aws, azure."
+        )

dataflow/secrets_manager/interface.py

@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+
+class SecretManager(ABC):
+    @abstractmethod
+    def create_secret(self, vault_path: str, secret_data: dict) -> str:
+        pass
+
+    @abstractmethod
+    def get_secret_by_key(self, vault_path: str):
+        pass
+
+    @abstractmethod
+    def update_secret(self, vault_path: str, update_data):
+        pass
+
+    @abstractmethod
+    def delete_secret(self, vault_path: str):
+        pass
+
+    @abstractmethod
+    def test_connection(self, vault_path: str):
+        pass