cornflow 2.0.0a13__tar.gz → 2.0.0a14__tar.gz

{cornflow-2.0.0a13/cornflow.egg-info → cornflow-2.0.0a14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cornflow
-Version: 2.0.0a13
+Version: 2.0.0a14
 Summary: cornflow is an open source multi-solver optimization server with a REST API built using flask.
 Home-page: https://github.com/baobabsoluciones/cornflow
 Author: baobab soluciones
@@ -14,7 +14,7 @@ Requires-Dist: alembic==1.9.2
 Requires-Dist: apispec<=6.3.0
 Requires-Dist: cachetools==5.3.3
 Requires-Dist: click<=8.1.7
-Requires-Dist: cornflow-client==2.0.0a13
+Requires-Dist: cornflow-client==2.0.0a14
 Requires-Dist: cryptography<=44.0.1
 Requires-Dist: databricks-sdk==0.29.0
 Requires-Dist: disposable-email-domains>=0.0.86

{cornflow-2.0.0a13 → cornflow-2.0.0a14}/cornflow/app.py

@@ -51,6 +51,8 @@ def create_app(env_name="development", dataconn=None):
     """
     dictConfig(log_config(app_config[env_name].LOG_LEVEL))
 
+    # Note: Explicit CSRF protection is not configured as the application uses
+    # JWT for authentication via headers, mitigating standard CSRF vulnerabilities.
     app = Flask(__name__)
     app.json.sort_keys = False
     app.logger.setLevel(app_config[env_name].LOG_LEVEL)
@@ -103,7 +105,7 @@ def create_app(env_name="development", dataconn=None):
     else:
         raise ConfigurationError(
             error="Invalid authentication type",
-            log_txt="Error while configuring authentication. The authentication type is not valid."
+            log_txt="Error while configuring authentication. The authentication type is not valid.",
         )
 
     initialize_errorhandlers(app)

{cornflow-2.0.0a13 → cornflow-2.0.0a14}/cornflow/cli/__init__.py

@@ -17,6 +17,10 @@ from cornflow.cli.views import views
 
 @click.group(name="cornflow", help="Commands in the cornflow cli")
 def cli():
+    """
+    This method is empty but it serves as the building block
+    for the rest of the commands
+    """
     pass
 
 

{cornflow-2.0.0a13 → cornflow-2.0.0a14}/cornflow/cli/actions.py

@@ -7,6 +7,10 @@ from .arguments import verbose
 
 @click.group(name="actions", help="Commands to manage the actions")
 def actions():
+    """
+    This method is empty but it serves as the building block
+    for the rest of the commands
+    """
     pass
 
 

{cornflow-2.0.0a13 → cornflow-2.0.0a14}/cornflow/cli/config.py

@@ -7,6 +7,10 @@ from .arguments import path
 
 @click.group(name="config", help="Commands to manage the configuration variables")
 def config():
+    """
+    This method is empty but it serves as the building block
+    for the rest of the commands
+    """
     pass
 
 

{cornflow-2.0.0a13 → cornflow-2.0.0a14}/cornflow/cli/migrations.py

@@ -3,6 +3,7 @@ import os.path
 
 import click
 from cornflow.shared import db
+from cornflow.shared.const import MIGRATIONS_DEFAULT_PATH
 from flask_migrate import Migrate, migrate, upgrade, downgrade, init
 
 from .utils import get_app
@@ -10,6 +11,10 @@ from .utils import get_app
 
 @click.group(name="migrations", help="Commands to manage the migrations")
 def migrations():
+    """
+    This method is empty but it serves as the building block
+    for the rest of the commands
+    """
     pass
 
 
@@ -18,12 +23,12 @@ def migrate_migrations():
     app = get_app()
     external = int(os.getenv("EXTERNAL_APP", 0))
     if external == 0:
-        path = "./cornflow/migrations"
+        path = MIGRATIONS_DEFAULT_PATH
     else:
         path = f"./{os.getenv('EXTERNAL_APP_MODULE', 'external_app')}/migrations"
 
     with app.app_context():
-        migration_client = Migrate(app=app, db=db, directory=path)
+        Migrate(app=app, db=db, directory=path)
         migrate()
 
 
@@ -35,12 +40,12 @@ def upgrade_migrations(revision="head"):
     app = get_app()
     external = int(os.getenv("EXTERNAL_APP", 0))
     if external == 0:
-        path = "./cornflow/migrations"
+        path = MIGRATIONS_DEFAULT_PATH
     else:
         path = f"./{os.getenv('EXTERNAL_APP_MODULE', 'external_app')}/migrations"
 
     with app.app_context():
-        migration_client = Migrate(app=app, db=db, directory=path)
+        Migrate(app=app, db=db, directory=path)
         upgrade(revision=revision)
 
 
@@ -52,12 +57,12 @@ def downgrade_migrations(revision="-1"):
     app = get_app()
     external = int(os.getenv("EXTERNAL_APP", 0))
     if external == 0:
-        path = "./cornflow/migrations"
+        path = MIGRATIONS_DEFAULT_PATH
     else:
         path = f"./{os.getenv('EXTERNAL_APP_MODULE', 'external_app')}/migrations"
 
     with app.app_context():
-        migration_client = Migrate(app=app, db=db, directory=path)
+        Migrate(app=app, db=db, directory=path)
         downgrade(revision=revision)
 
 
@@ -69,10 +74,10 @@ def init_migrations():
     app = get_app()
     external = int(os.getenv("EXTERNAL_APP", 0))
     if external == 0:
-        path = "./cornflow/migrations"
+        path = MIGRATIONS_DEFAULT_PATH
     else:
         path = f"./{os.getenv('EXTERNAL_APP_MODULE', 'external_app')}/migrations"
 
     with app.app_context():
-        migration_client = Migrate(app=app, db=db, directory=path)
+        Migrate(app=app, db=db, directory=path)
         init()

{cornflow-2.0.0a13 → cornflow-2.0.0a14}/cornflow/cli/permissions.py

@@ -11,6 +11,10 @@ from .utils import get_app
 
 @click.group(name="permissions", help="Commands to manage the permissions")
 def permissions():
+    """
+    This method is empty but it serves as the building block
+    for the rest of the commands
+    """
     pass
 
 

{cornflow-2.0.0a13 → cornflow-2.0.0a14}/cornflow/cli/roles.py

@@ -6,6 +6,10 @@ from .utils import get_app
 
 @click.group(name="roles", help="Commands to manage the roles")
 def roles():
+    """
+    This method is empty but it serves as the building block
+    for the rest of the commands
+    """
     pass
 
 

{cornflow-2.0.0a13 → cornflow-2.0.0a14}/cornflow/cli/schemas.py

@@ -1,6 +1,7 @@
 """
 File that implements the generate from schema cli command
 """
+
 import click
 from .tools.api_generator import APIGenerator
 from .tools.schema_generator import SchemaGenerator
@@ -20,6 +21,10 @@ METHOD_OPTIONS = [
 
 @click.group(name="schemas", help="Commands to manage the schemas")
 def schemas():
+    """
+    This method is empty but it serves as the building block
+    for the rest of the commands
+    """
     pass
 
 

cornflow-2.0.0a14/cornflow/cli/service.py

@@ -0,0 +1,397 @@
+import os
+import subprocess
+import sys
+import time
+from logging import error
+
+
+import click
+from .utils import get_db_conn
+import cornflow
+from cornflow.app import create_app
+from cornflow.commands import (
+    access_init_command,
+    create_user_with_role,
+    register_deployed_dags_command,
+    register_dag_permissions_command,
+    update_schemas_command,
+    update_dag_registry_command,
+)
+from cornflow.shared.const import (
+    AUTH_DB,
+    AUTH_LDAP,
+    AUTH_OID,
+    ADMIN_ROLE,
+    SERVICE_ROLE,
+    PLANNER_ROLE,
+    DATABRICKS_BACKEND,
+    AIRFLOW_BACKEND,
+)
+from cornflow.shared import db
+from cryptography.fernet import Fernet
+from flask_migrate import Migrate, upgrade
+
+MAIN_WD = "/usr/src/app"
+
+
+@click.group(name="service", help="Commands to run the cornflow service")
+def service():
+    """
+    This method is empty but it serves as the building block
+    for the rest of the commands
+    """
+    pass
+
+
+@service.command(name="init", help="Initialize the service")
+def init_cornflow_service():
+    click.echo("Starting the service")
+    os.chdir(MAIN_WD)
+
+    config = _setup_environment_variables()
+    _configure_logging(config["cornflow_logging"])
+
+    external_application = config["external_application"]
+    environment = config["environment"]
+    cornflow_db_conn = config["cornflow_db_conn"]
+    external_app_module = config["external_app_module"]
+
+    app = None  # Initialize app to None
+
+    if external_application == 0:
+        click.echo("Initializing standard Cornflow application")
+        app = create_app(environment, cornflow_db_conn)
+        with app.app_context():
+            _initialize_database(app)
+            _create_initial_users(
+                config["auth"],
+                config["cornflow_admin_user"],
+                config["cornflow_admin_email"],
+                config["cornflow_admin_pwd"],
+                config["cornflow_service_user"],
+                config["cornflow_service_email"],
+                config["cornflow_service_pwd"],
+            )
+            _sync_with_airflow(
+                config["airflow_url"],
+                config["airflow_user"],
+                config["airflow_pwd"],
+                config["open_deployment"],
+                external_app=False,
+            )
+        _start_application(external_application, environment)
+
+    elif external_application == 1:
+        click.echo(f"Initializing Cornflow with external app: {external_app_module}")
+        if not external_app_module:
+            sys.exit("FATAL: EXTERNAL_APP is 1 but EXTERNAL_APP_MODULE is not set.")
+
+        _setup_external_app()
+        from importlib import import_module
+
+        external_app_lib = import_module(external_app_module)
+        app = external_app_lib.create_wsgi_app(environment, cornflow_db_conn)
+
+        with app.app_context():
+            _initialize_database(app, external_app_module)
+            _create_initial_users(
+                config["auth"],
+                config["cornflow_admin_user"],
+                config["cornflow_admin_email"],
+                config["cornflow_admin_pwd"],
+                config["cornflow_service_user"],
+                config["cornflow_service_email"],
+                config["cornflow_service_pwd"],
+            )
+            _sync_with_airflow(
+                config["airflow_url"],
+                config["airflow_user"],
+                config["airflow_pwd"],
+                config["open_deployment"],
+                external_app=True,
+            )
+        _start_application(external_application, environment, external_app_module)
+
+    else:
+        # This case should ideally be caught earlier or handled differently
+        sys.exit(f"FATAL: Invalid EXTERNAL_APP value: {external_application}")
+
+
+def _setup_environment_variables():
+    """Reads environment variables, sets defaults, and returns config values."""
+    environment = os.getenv("FLASK_ENV", "development")
+    os.environ["FLASK_ENV"] = environment
+
+    ###################################
+    # Global defaults and back-compat #
+    ###################################
+    # cornflow backend selection
+    cornflow_backend = os.getenv("CORNFLOW_BACKEND", str(AIRFLOW_BACKEND))
+    os.environ["CORNFLOW_BACKEND"] = cornflow_backend
+    cornflow_backend = int(cornflow_backend)
+    # Airflow global default conn
+    if cornflow_backend == AIRFLOW_BACKEND:
+        airflow_user = os.getenv("AIRFLOW_USER", "admin")
+        airflow_pwd = os.getenv("AIRFLOW_PWD", "admin")
+        airflow_url = os.getenv("AIRFLOW_URL", "http://webserver:8080")
+        os.environ["AIRFLOW_USER"] = airflow_user
+        os.environ["AIRFLOW_PWD"] = airflow_pwd
+        os.environ["AIRFLOW_URL"] = airflow_url
+    elif cornflow_backend == DATABRICKS_BACKEND:
+        databricks_url = os.getenv("DATABRICKS_HOST")
+        databricks_auth_secret = os.getenv("DATABRICKS_CLIENT_SECRET")
+        databricks_token_endpoint = os.getenv("DATABRICKS_TOKEN_ENDPOINT")
+        databricks_ep_clusters = os.getenv("DATABRICKS_EP_CLUSTERS")
+        databricks_client_id = os.getenv("DATABRICKS_CLIENT_ID")
+        os.environ["DATABRICKS_HOST"] = databricks_url
+        os.environ["DATABRICKS_CLIENT_SECRET"] = databricks_auth_secret
+        os.environ["DATABRICKS_TOKEN_ENDPOINT"] = databricks_token_endpoint
+        os.environ["DATABRICKS_EP_CLUSTERS"] = databricks_ep_clusters
+        os.environ["DATABRICKS_CLIENT_ID"] = databricks_client_id
+    else:
+        raise Exception("Selected backend not among valid options")
+    # Cornflow app config
+    os.environ.setdefault("cornflow_url", "http://cornflow:5000")
+    os.environ["FLASK_APP"] = "cornflow.app"
+    os.environ["SECRET_KEY"] = os.getenv("FERNET_KEY", Fernet.generate_key().decode())
+
+    # Cornflow db defaults
+    os.environ["DEFAULT_POSTGRES"] = "1"
+    cornflow_db_conn = get_db_conn()
+    os.environ["DATABASE_URL"] = cornflow_db_conn
+
+    # Platform auth config and service users
+    auth = int(os.getenv("AUTH_TYPE", AUTH_DB))
+    cornflow_admin_user = os.getenv("CORNFLOW_ADMIN_USER", "cornflow_admin")
+    cornflow_admin_email = os.getenv(
+        "CORNFLOW_ADMIN_EMAIL", "cornflow_admin@cornflow.com"
+    )
+    cornflow_admin_pwd = os.getenv("CORNFLOW_ADMIN_PWD", "Cornflow_admin1234")
+    cornflow_service_user = os.getenv("CORNFLOW_SERVICE_USER", "service_user")
+    cornflow_service_email = os.getenv(
+        "CORNFLOW_SERVICE_EMAIL", "service_user@cornflow.com"
+    )
+    cornflow_service_pwd = os.getenv("CORNFLOW_SERVICE_PWD", "Service_user1234")
+
+    # Cornflow logging and deployment config
+    cornflow_logging = os.getenv("CORNFLOW_LOGGING", "console")
+    os.environ["CORNFLOW_LOGGING"] = cornflow_logging
+    open_deployment = os.getenv("OPEN_DEPLOYMENT", 1)
+    os.environ["OPEN_DEPLOYMENT"] = str(open_deployment)
+    signup_activated = os.getenv("SIGNUP_ACTIVATED", 1)
+    os.environ["SIGNUP_ACTIVATED"] = str(signup_activated)
+    user_access_all_objects = os.getenv("USER_ACCESS_ALL_OBJECTS", 0)
+    os.environ["USER_ACCESS_ALL_OBJECTS"] = str(user_access_all_objects)
+    default_role = int(os.getenv("DEFAULT_ROLE", PLANNER_ROLE))
+    os.environ["DEFAULT_ROLE"] = str(default_role)
+
+    # Check LDAP parameters for active directory and show message
+    if auth == AUTH_LDAP:
+        click.echo(
+            "WARNING: Cornflow will be deployed with LDAP Authorization. "
+            "Please review your ldap auth configuration."
+        )
+
+    # check database param from docker env
+    if cornflow_db_conn is None:
+        sys.exit("FATAL: you need to provide a postgres database for Cornflow")
+
+    external_application = int(os.getenv("EXTERNAL_APP", 0))
+    external_app_module = os.getenv("EXTERNAL_APP_MODULE")
+    if cornflow_backend == AIRFLOW_BACKEND:
+        return {
+            "environment": environment,
+            "auth": auth,
+            "airflow_user": airflow_user,
+            "airflow_pwd": airflow_pwd,
+            "airflow_url": airflow_url,
+            "cornflow_db_conn": cornflow_db_conn,
+            "cornflow_admin_user": cornflow_admin_user,
+            "cornflow_admin_email": cornflow_admin_email,
+            "cornflow_admin_pwd": cornflow_admin_pwd,
+            "cornflow_service_user": cornflow_service_user,
+            "cornflow_service_email": cornflow_service_email,
+            "cornflow_service_pwd": cornflow_service_pwd,
+            "cornflow_logging": cornflow_logging,
+            "open_deployment": open_deployment,
+            "external_application": external_application,
+            "external_app_module": external_app_module,
+        }
+    elif cornflow_backend == DATABRICKS_BACKEND:
+        return {
+            "environment": environment,
+            "auth": auth,
+            "databricks_url": databricks_url,
+            "databricks_auth_secret": databricks_auth_secret,
+            "databricks_token_endpoint": databricks_token_endpoint,
+            "databricks_ep_clusters": databricks_ep_clusters,
+            "databricks_client_id": databricks_client_id,
+            "cornflow_db_conn": cornflow_db_conn,
+            "cornflow_admin_user": cornflow_admin_user,
+            "cornflow_admin_email": cornflow_admin_email,
+            "cornflow_admin_pwd": cornflow_admin_pwd,
+            "cornflow_service_user": cornflow_service_user,
+            "cornflow_service_email": cornflow_service_email,
+            "cornflow_service_pwd": cornflow_service_pwd,
+            "cornflow_logging": cornflow_logging,
+            "open_deployment": open_deployment,
+            "external_application": external_application,
+            "external_app_module": external_app_module,
+        }
+    else:
+        raise Exception("Selected backend not among valid options")
+
+
+def _configure_logging(cornflow_logging):
+    """Configures log rotation if logging to file."""
+    if cornflow_logging == "file":
+        try:
+            conf = f"""/usr/src/app/log/*.log {{
+            rotate 30
+            daily
+            compress
+            size 20M
+            postrotate
+             kill -HUP $(cat {MAIN_WD}/gunicorn.pid)
+            endscript}}"""
+            logrotate = subprocess.run(
+                f"cat > /etc/logrotate.d/cornflow <<EOF\n {conf} \nEOF",
+                shell=True,
+                capture_output=True,
+                text=True,
+            )
+            if logrotate.returncode != 0:
+                error(f"Error configuring logrotate: {logrotate.stderr}")
+            else:
+                print(logrotate.stdout)
+        except Exception as e:
+            error(f"Exception during logrotate configuration: {e}")
+
+
+def _initialize_database(app, external_app_module=None):
+    """Initializes the database and runs migrations."""
+    with app.app_context():
+        if external_app_module:
+            from importlib import import_module
+
+            external_app_lib = import_module(external_app_module)
+            migrations_path = f"{os.path.dirname(external_app_lib.__file__)}/migrations"
+        else:
+            migrations_path = f"{os.path.dirname(cornflow.__file__)}/migrations"
+
+        Migrate(app=app, db=db, directory=migrations_path)
+        upgrade()
+        access_init_command(verbose=False)
+
+
+def _create_initial_users(
+    auth,
+    admin_user,
+    admin_email,
+    admin_pwd,
+    service_user,
+    service_email,
+    service_pwd,
+):
+    """Creates the initial admin and service users if using DB or OID auth."""
+    if auth == AUTH_DB or auth == AUTH_OID:
+        # create cornflow admin user
+        create_user_with_role(
+            admin_user,
+            admin_email,
+            admin_pwd,
+            "admin",
+            ADMIN_ROLE,
+            verbose=True,
+        )
+        # create cornflow service user
+        create_user_with_role(
+            service_user,
+            service_email,
+            service_pwd,
+            "serviceuser",
+            SERVICE_ROLE,
+            verbose=True,
+        )
+
+
+def _sync_with_airflow(
+    airflow_url, airflow_user, airflow_pwd, open_deployment, external_app=False
+):
+    """Syncs DAGs, permissions, and schemas with Airflow."""
+    register_deployed_dags_command(airflow_url, airflow_user, airflow_pwd, verbose=True)
+    register_dag_permissions_command(open_deployment, verbose=True)
+    update_schemas_command(airflow_url, airflow_user, airflow_pwd, verbose=True)
+    if external_app:
+        update_dag_registry_command(
+            airflow_url, airflow_user, airflow_pwd, verbose=True
+        )
+
+
+def _setup_external_app():
+    """Performs setup steps specific to external applications."""
+    os.chdir(MAIN_WD)
+    if _register_key():
+        prefix = "CUSTOM_SSH_"
+        env_variables = {
+            key: value for key, value in os.environ.items() if key.startswith(prefix)
+        }
+        for _, value in env_variables.items():
+            _register_ssh_host(value)
+
+    # Install requirements for the external app
+    pip_install_cmd = "$(command -v pip) install --user -r requirements.txt"
+    click.echo(f"Running: {pip_install_cmd}")
+    result = subprocess.run(pip_install_cmd, shell=True, capture_output=True, text=True)
+    if result.returncode != 0:
+        error(f"Error installing requirements: {result.stderr}")
+    else:
+        print(result.stdout)
+    time.sleep(5)  # Consider if this sleep is truly necessary
+    sys.path.append(MAIN_WD)
+
+
+def _start_application(external_application, environment, external_app_module=None):
+    """Starts the Gunicorn server."""
+    if external_application == 0:
+        os.environ["GUNICORN_WORKING_DIR"] = MAIN_WD
+        gunicorn_cmd = (
+            "/usr/local/bin/gunicorn -c python:cornflow.gunicorn "
+            f"\"cornflow.app:create_app('{environment}')\""
+        )
+    elif external_application == 1:
+        os.environ["GUNICORN_WORKING_DIR"] = MAIN_WD
+        if not external_app_module:
+            raise ValueError(
+                "EXTERNAL_APP_MODULE must be set for external applications"
+            )
+        gunicorn_cmd = (
+            "/usr/local/bin/gunicorn -c python:cornflow.gunicorn "
+            f"\"{external_app_module}:create_wsgi_app('{environment}')\""
+        )
+    else:
+        raise ValueError(f"Invalid EXTERNAL_APP value: {external_application}")
+
+    click.echo(f"Starting application with Gunicorn: {gunicorn_cmd}")
+    os.system(gunicorn_cmd)
+
+
+def _register_ssh_host(host):
+    if host is not None:
+        add_host = f"ssh-keyscan {host} >> {MAIN_WD}/.ssh/known_hosts"
+        config_ssh_host = f"echo Host {host} >> {MAIN_WD}/.ssh/config"
+        config_ssh_key = (
+            'echo "   IdentityFile {MAIN_WD}/.ssh/id_rsa" >> {MAIN_WD}/.ssh/config'
+        )
+        os.system(add_host)
+        os.system(config_ssh_host)
+        os.system(config_ssh_key)
+
+
+def _register_key():
+    if os.path.isfile(f"{MAIN_WD}/.ssh/id_rsa"):
+        add_key = f"chmod 0600 {MAIN_WD}/.ssh/id_rsa && ssh-add {MAIN_WD}/.ssh/id_rsa"
+        os.system(add_key)
+        return True
+    else:
+        return False

{cornflow-2.0.0a13 → cornflow-2.0.0a14}/cornflow/cli/tools/api_generator.py

@@ -1,6 +1,7 @@
 """
 This file has the class that creates the new API
 """
+
 import json
 import os
 import re
@@ -10,6 +11,8 @@ from .models_tools import ModelGenerator, model_shared_imports
 from .schemas_tools import SchemaGenerator, schemas_imports
 from .tools import generate_class_def
 
+INIT_FILE = "__init__.py"
+
 
 class APIGenerator:
     """
@@ -52,7 +55,7 @@ class APIGenerator:
         self.endpoint_path = os.path.join(self.output_path, "endpoints")
         self.schema_path = os.path.join(self.output_path, "schemas")
         self.init_resources = []
-        self.init_file = os.path.join(self.endpoint_path, "__init__.py")
+        self.init_file = os.path.join(self.endpoint_path, INIT_FILE)
 
     def import_schema(self) -> dict:
         """
@@ -77,23 +80,23 @@ class APIGenerator:
         if not os.path.isdir(self.model_path):
             os.mkdir(self.model_path)
 
-        init_path = os.path.join(self.model_path, "__init__.py")
+        init_path = os.path.join(self.model_path, INIT_FILE)
         with open(init_path, "w") as file:
-            file.write(f'"""\nThis file exposes the models\n"""\n')
+            file.write("\nThis file exposes the models\n\n")
 
         if not os.path.isdir(self.endpoint_path):
             os.mkdir(self.endpoint_path)
 
-        init_path = os.path.join(self.endpoint_path, "__init__.py")
+        init_path = os.path.join(self.endpoint_path, INIT_FILE)
         with open(init_path, "w") as file:
-            file.write(f'"""\nThis file exposes the endpoints\n"""\n')
+            file.write("\nThis file exposes the endpoints\n\n")
 
         if not os.path.isdir(self.schema_path):
             os.mkdir(self.schema_path)
 
-        init_path = os.path.join(self.schema_path, "__init__.py")
+        init_path = os.path.join(self.schema_path, INIT_FILE)
         with open(init_path, "w") as file:
-            file.write(f'"""\nThis file exposes the schemas\n"""\n')
+            file.write("\nThis file exposes the schemas\n\n")
 
     def main(self):
         """
@@ -151,7 +154,7 @@ class APIGenerator:
             fd.write(mg.generate_model_repr_str())
             fd.write("\n")
 
-        init_file = os.path.join(self.model_path, "__init__.py")
+        init_file = os.path.join(self.model_path, INIT_FILE)
 
         with open(init_file, "a") as file:
             file.write(f"from .{self.prefix}{table_name} import {class_name}\n")
@@ -207,7 +210,7 @@ class APIGenerator:
             fd.write(generate_class_def(class_name_one, parents_class))
             fd.write(sg.generate_schema())
 
-        init_file = os.path.join(self.schema_path, "__init__.py")
+        init_file = os.path.join(self.schema_path, INIT_FILE)
         with open(init_file, "a") as file:
             file.write(
                 f"from .{self.prefix}{table_name} import {class_name_one}, "
@@ -432,7 +435,7 @@ class APIGenerator:
         :return: str: the type in format "<type:idx>"
         """
         schema_table = self.schema["properties"][table_name]["items"]["properties"]
-        id_type=None
+        id_type = None
         if "id" in schema_table.keys():
             id_type = schema_table["id"]["type"]
         if id_type == "string" or isinstance(id_type, list):