dataflow-core 2.1.15rc1__tar.gz → 2.1.15rc3__tar.gz

{dataflow_core-2.1.15rc1 → dataflow_core-2.1.15rc3}/PKG-INFO

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

{dataflow_core-2.1.15rc1 → dataflow_core-2.1.15rc3}/authenticator/dataflowhubauthenticator.py

@@ -8,7 +8,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 +29,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 +45,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,23 +71,6 @@ 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:
@@ -103,16 +80,12 @@ class DataflowBaseAuthenticator(Authenticator):
     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 +95,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)
@@ -135,8 +108,6 @@ class DataflowBaseAuthenticator(Authenticator):
                 first_name=first_name or username,
                 last_name=last_name or "",
                 email=email,
-                role_id=role_id,
-                password='user@123',
             )
             
             self.db.add(new_user)
@@ -154,22 +125,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_core-2.1.15rc1 → dataflow_core-2.1.15rc3}/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_core-2.1.15rc1 → dataflow_core-2.1.15rc3}/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_core-2.1.15rc1 → dataflow_core-2.1.15rc3}/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_core-2.1.15rc3/dataflow/models/app_types.py

@@ -0,0 +1,15 @@
+from sqlalchemy import Column, Integer, String, Boolean
+from sqlalchemy.orm import relationship
+from dataflow.db import Base
+
+class AppType(Base):
+    __tablename__ = "APP_TYPE"
+    
+    id = Column(Integer, primary_key=True, autoincrement=True, unique=True)
+    name = Column(String, unique=True, nullable=False)
+    display_name = Column(String, 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_core-2.1.15rc1 → dataflow_core-2.1.15rc3}/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_core-2.1.15rc1 → dataflow_core-2.1.15rc3}/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")
 

dataflow_core-2.1.15rc3/dataflow/models/environment.py

@@ -0,0 +1,125 @@
+from sqlalchemy import (
+    Column, Integer, String, Boolean, Text, 
+    ForeignKey, DateTime, UniqueConstraint, CheckConstraint
+)
+from sqlalchemy.orm import relationship, Session
+from sqlalchemy.sql import func
+from datetime import datetime, timezone
+from dataflow.db import Base
+from enum import Enum
+
+class EnvironmentAttributes(Base):
+    """
+    Shared columns between Environment and ArchivedEnvironment.
+    """
+    __abstract__ = True 
+
+    name = Column(String, nullable=False)
+    url = Column(String)
+    enabled = Column(Boolean, default=True, server_default='true')
+    version = Column(String, default=0, server_default='0')
+    is_latest = Column(Boolean, default=True, server_default='true')
+    base_env_id = Column(Integer, default=None)
+    short_name = Column(String(5))
+    status = Column(String, default="Saved", server_default="Saved")
+    icon = Column(String)
+    py_version = Column(String)
+    r_version = Column(String)
+    pip_libraries = Column(Text)
+    conda_libraries = Column(Text)
+    r_requirements = Column(Text)
+    created_date = Column(DateTime, server_default=func.now())
+    created_by = Column(String)
+    org_id = Column(Integer, ForeignKey('ORGANIZATION.id'))
+
+class Environment(EnvironmentAttributes): 
+    __tablename__ = 'ENVIRONMENT'
+    __table_args__ = (UniqueConstraint('short_name', 'org_id', name='_env_short_name_org_uc'),)
+    id = Column(Integer, primary_key=True, autoincrement=True)
+
+    # Relationships
+    organization = relationship("Organization", back_populates="environments")
+    archived_versions = relationship("ArchivedEnvironment", back_populates="original_environment")
+
+class ArchivedEnvironment(EnvironmentAttributes):
+    __tablename__ = 'ARCHIVED_ENVIRONMENT'
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    original_env_id = Column(Integer, ForeignKey('ENVIRONMENT.id', ondelete='CASCADE'))
+
+    # Relationship with Environment
+    original_environment = relationship("Environment", back_populates="archived_versions")
+
+class JobLogs(Base):
+    __tablename__ = "JOB_LOG"
+    __table_args__ = (UniqueConstraint('log_file_name', 'org_id', name='_job_log_file_org_uc'),)
+
+    id = Column(Integer, primary_key=True, index=True)
+    created_at = Column(DateTime, default=datetime.now, server_default=func.now())
+    completed_at = Column(DateTime, nullable=True)
+    log_file_name = Column(String, nullable=False)
+    log_file_location = Column(String, nullable=False)
+    status = Column(String)
+    created_by = Column(String)
+    org_id = Column(Integer, ForeignKey('ORGANIZATION.id', ondelete='CASCADE'))
+
+
+class LocalEnvironment(Base):
+    __tablename__ = "LOCAL_ENVIRONMENT"
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    name = Column(String, nullable=False, index=True)
+    user_name = Column(String, ForeignKey('USER.user_name', ondelete='CASCADE'), nullable=False, index=True)
+    org_id = Column(Integer, ForeignKey('ORGANIZATION.id', ondelete='CASCADE'), nullable=False, index=True)
+    py_version = Column(String)
+    pip_libraries = Column(Text)
+    conda_libraries = Column(Text)
+    status = Column(String, default="Created", server_default="Created")
+    cloned_from = Column(String, nullable=True)
+    updated_at = Column(DateTime, default=datetime.now(timezone.utc), onupdate=datetime.now(timezone.utc))
+    need_refresh = Column(Boolean, default=False, server_default='false')
+
+class EnvType(str, Enum):
+    dataflow = "dataflow"
+    local = "local"
+
+class PipSource(Base):
+    __tablename__ = "PIP_SOURCE"
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+
+    org_id = Column(Integer, ForeignKey("ORGANIZATION.id", ondelete="CASCADE"), nullable=False, index=True)
+    user_name = Column(String, ForeignKey("USER.user_name", ondelete="CASCADE"), nullable=True, index=True)
+
+    name = Column(String, nullable=False)
+    url = Column(String, nullable=False)
+    is_index = Column(Boolean, default=False, nullable=False, server_default='false')
+
+    created_at = Column(DateTime, default=datetime.now(timezone.utc), nullable=False)
+    updated_at = Column(DateTime, default=datetime.now(timezone.utc), onupdate=datetime.now(timezone.utc), nullable=False)
+
+    __table_args__ = (
+        UniqueConstraint("org_id", "name", "user_name", name="uq_pip_source_per_user_org"),
+        CheckConstraint("NOT (is_index = TRUE AND user_name IS NOT NULL)", name="check_no_user_index_url"),
+    )
+
+    @classmethod
+    def get_org_sources(cls, session: Session, org_id: int):
+        """
+        Returns all sources for the given org (org-level).
+        """
+        return session.query(cls).filter(
+            cls.org_id == org_id,
+            cls.user_name == None
+        ).all()
+
+    @classmethod
+    def get_user_sources(cls, session: Session, org_id: int, user_name: str):
+        """
+        Returns merged sources for a user in an org (org-level + user-level personal sources).
+        """
+        return session.query(cls).filter(
+            cls.org_id == org_id,
+            ((cls.user_name == None) | (cls.user_name == user_name))
+        ).all()
+

{dataflow_core-2.1.15rc1 → dataflow_core-2.1.15rc3}/dataflow/models/git_ssh.py

@@ -8,11 +8,12 @@ class GitSSH(Base):
 
     id = Column(Integer, primary_key=True, index=True, autoincrement=True)
     user_name = Column(String, ForeignKey('USER.user_name', ondelete="CASCADE"), nullable=False)
+    org_id = Column(Integer, ForeignKey("ORGANIZATION.id"), index=True, nullable=False)
     description = Column(String)
     key_name = Column(String, nullable=False)
     created_date = Column(DateTime, server_default=func.now(), nullable=False)
     last_used_date = Column(DateTime)
 
     __table_args__ = (
-        UniqueConstraint(user_name, key_name, name='user_name_key_name_unique'),
+        UniqueConstraint(user_name, key_name, org_id, name='user_name_key_name_org_id_unique'),
     )