[Datastore] Fix SQL injection vulnerabilities (#3680)

mlrun · May 31, 2023 · 3532025 · 3532025
1 parent 17cdc1d
commit 3532025
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/mlrun/datastore/sources.py b/mlrun/datastore/sources.py
@@ -943,14 +943,17 @@ def to_dataframe(self):
         query = self.attributes.get("query", None)
         db_path = self.attributes.get("db_path")
         table_name = self.attributes.get("table_name")
+        params = None
         if not query:
-            query = f"SELECT * FROM {table_name}"
+            query = "SELECT * FROM %(table)s"
+            params = {"table": table_name}
         if table_name and db_path:
             engine = db.create_engine(db_path)
             with engine.connect() as con:
                 return pd.read_sql(
                     query,
                     con=con,
+                    params=params,
                     chunksize=self.attributes.get("chunksize"),
                     parse_dates=self.attributes.get("time_fields"),
                 )

diff --git a/mlrun/datastore/targets.py b/mlrun/datastore/targets.py
@@ -1645,8 +1645,9 @@ def as_df(
         engine = sqlalchemy.create_engine(db_path)
         with engine.connect() as conn:
             df = pd.read_sql(
-                f"SELECT * FROM {self.attributes.get('table_name')}",
+                "SELECT * FROM %(table)s",
                 con=conn,
+                params={"table": self.attributes.get("table_name")},
                 parse_dates=self.attributes.get("time_fields"),
             )
             if self._primary_key_column: