Apply offset only if not null

cornflow-server/cornflow/cli/tools/schema_generator.py

@@ -26,11 +26,13 @@ def __init__(self, path, output_path=None, ignore_files=None, leave_bases=False)
         self.data = SuperDict()
         self.model_table = dict()
         self.table_model = dict()
+        self.prefix_table_name = "sfm_tmp_"
 
     def main(self):
+        self.clear()
         os.mkdir(self.tmp_path)
 
-        copy_tree(self.path, self.tmp_path)
+        self.copy_tree(self.path, self.tmp_path)
 
         files = (
             TupList(os.listdir(self.tmp_path))
@@ -43,19 +45,30 @@ def main(self):
             .vapply(lambda v: (os.path.join(self.tmp_path, v), v[:-3]))
         )
 
+        self.replace_table_names(files)
+
         self.mock_packages(files)
 
         self.parse(files)
 
         self.inherit()
 
+        self.clean_table_names()
+
         schema = self.to_schema()
 
         with open(self.output_path, "w") as fd:
             json.dump(schema, fd, indent=2)
         self.clear()
         return 0
 
+    def copy_tree(self, origin_path, destination_path):
+        # Creating a temporary copy of all python files
+        for root, dirs, files in os.walk(origin_path):
+            for file_ in files:
+                if file_.endswith(".py"):
+                    shutil.copy(os.path.join(root, file_), os.path.join(destination_path, file_))
+
     def mock_packages(self, files):
         # Mocking all relative imports
         for file_path, file_name in files:
@@ -67,11 +80,47 @@ def mock_packages(self, files):
             libs = re.findall(r"from ((\.+\w+)+) import (\w+)", text)
             for lib in libs:
                 text = text.replace(lib[0], "mockedpackage")
+            text = text.replace("@declared_attr", "")
             with open(file_path, "w") as fd:
                 fd.write(text)
 
         sys.modules["mockedpackage"] = MagicMock()
 
+    def replace_table_names(self, files):
+        # Mock tablenames
+        for file_path, file_name in files:
+            with open(file_path, "r") as fd:
+                text = fd.read()
+            tablenames_1 = list(re.findall(r"__tablename__ = \"(.+)\"", text)) + list(re.findall(r"__tablename__ = \'(.+)\'", text))
+            tablenames_2 = list(re.findall(r"__tablename__ =\"(.+)\"", text)) + list(re.findall(r"__tablename__ =\'(.+)\'", text))
+            for tab in tablenames_1:
+                text = text.replace(f"__tablename__ = \"{tab}\"", f"__tablename__ = '{self.prefix_table_name}{tab}'")
+                text = text.replace(f"__tablename__ = \'{tab}\'", f"__tablename__ = '{self.prefix_table_name}{tab}'")
+            for tab in tablenames_2:
+                text = text.replace(f"__tablename__ =\"{tab}\"", f"__tablename__ = '{self.prefix_table_name}{tab}'")
+                text = text.replace(f"__tablename__ =\'{tab}\'", f"__tablename__ = '{self.prefix_table_name}{tab}'")
+            parents = re.findall(r"class (.+)\((.+)\):", text)
+            for cl, parent in parents:
+                if parent != "db.Model":
+                    new_parent = self.prefix_table_name + parent
+                    text = text.replace(
+                        f"class {cl}({parent})",
+                        f"class {self.prefix_table_name + cl}({new_parent})"
+                    )
+                    imports = re.findall(r"from (\w*(\.+\w+)*)" + f" import {parent}", text)
+                    for imp in imports:
+                        text = text.replace(
+                            f"from {imp[0]} import {parent}",
+                            f"from mockedpackage import {new_parent}"
+                        )
+                else:
+                    text = text.replace(
+                        f"class {cl}({parent})",
+                        f"class {self.prefix_table_name + cl}({parent})"
+                    )
+            with open(file_path, "w") as fd:
+                fd.write(text)
+
     def parse(self, files):
         forget_keys = ["created_at", "updated_at", "deleted_at"]
         db = SQLAlchemy()
@@ -109,6 +158,15 @@ def parse(self, files):
                     self.model_table[model] = table_name
                     self.table_model[table_name] = model
                     for key, val in props.items():
+                        if callable(val):
+                            # declared_attr
+                            try:
+                                tmp_val = val(None)
+                                if isinstance(tmp_val, db.Column):
+                                    val = tmp_val
+                            except:
+                                pass
+
                         if key in forget_keys:
                             continue
                         elif isinstance(val, db.Column):
@@ -145,6 +203,14 @@ def parse(self, files):
             click.echo(err)
 
     def inherit(self):
+        self.parents = {
+            k.replace(self.prefix_table_name, ""): None if v is None else v.replace(self.prefix_table_name, "")
+            for k, v in self.parents.items()
+        }
+        self.model_table = {
+            k.replace(self.prefix_table_name, ""): v
+            for k, v in self.model_table.items()
+        }
         all_classes = set(self.parents.keys())
         not_treated = set(all_classes)
         treated = {"db.Model"}
@@ -174,6 +240,12 @@ def inherit(self):
         if not self.leave_bases:
             self.data = self.data.vfilter(lambda v: not v.get("remove", False))
 
+    def clean_table_names(self):
+        self.data = {
+            k.replace(self.prefix_table_name, ""): v
+            for k, v in self.data.items()
+        }
+
     def clear(self):
         if os.path.isdir(self.tmp_path):
             shutil.rmtree(self.tmp_path)

cornflow-server/cornflow/models/meta_models.py

@@ -154,7 +154,9 @@ def get_all_objects(
         """
         if "user" in kwargs:
             kwargs.pop("user")
-        query = cls.query.filter_by(**kwargs).offset(offset)
+        query = cls.query.filter_by(**kwargs)
+        if offset:
+            query = query.offset(offset)
         if limit:
             query = query.limit(limit)
         return query
@@ -300,7 +302,8 @@ def get_all_objects(
         if creation_date_lte:
             query = query.filter(cls.created_at <= creation_date_lte)
 
-        query = query.offset(offset)
+        if offset:
+            query = query.offset(offset)
         if limit:
             query = query.limit(limit)
         return query