update code formatting - use black

Author: ori-n

src/sql_data_guard/rest/sql_data_guard_rest.py

@@ -8,7 +8,8 @@
 
 app = Flask(__name__)
 
-@app.route('/verify-sql', methods=['POST'])
+
+@app.route("/verify-sql", methods=["POST"])
 def _verify_sql():
     if not request.is_json:
         return jsonify({"error": "Request must be JSON"}), 400
@@ -24,15 +25,15 @@ def _verify_sql():
     result["errors"] = list(result["errors"])
     return jsonify(result)
 
+
 def _init_logging():
     fileConfig(os.path.join(os.path.dirname(os.path.abspath(__file__)), "logging.conf"))
     logging.info("Logging initialized")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     _init_logging()
-    logging.getLogger("werkzeug").setLevel('WARNING')
+    logging.getLogger("werkzeug").setLevel("WARNING")
     port = os.environ.get("APP_PORT", 5000)
     logging.info(f"Going to start the app. Port: {port}")
     app.run(host="0.0.0.0", port=port)
-

src/sql_data_guard/sql_data_guard.py

@@ -18,6 +18,7 @@ class _VerificationContext:
         _dynamic_tables (Set[str]): Set of dynamic tables found in the query, like sub select and WITH clauses.
         _dialect (str): The SQL dialect to use for parsing.
     """
+
     def __init__(self, config: dict, dialect: str):
         super().__init__()
         self._can_fix = True
@@ -50,7 +51,6 @@ def fixed(self) -> Optional[str]:
     def fixed(self, value: Optional[str]):
         self._fixed = value
 
-
     @property
     def config(self) -> dict:
         return self._config
@@ -95,31 +95,48 @@ def verify_sql(sql: str, config: dict, dialect: str = None) -> dict:
         if isinstance(parsed, expr.Command):
             result.add_error(f"{parsed.name} statement is not allowed", False, 0.9)
         elif isinstance(parsed, expr.Delete) or isinstance(parsed, expr.Insert):
-            result.add_error(f"{parsed.key.upper()} statement is not allowed", False, 0.9)
+            result.add_error(
+                f"{parsed.key.upper()} statement is not allowed", False, 0.9
+            )
         elif isinstance(parsed, expr.Query):
             _verify_query_statement(parsed, result)
         else:
             result.add_error("Could not find a query statement", False, 0.7)
     if result.can_fix and len(result.errors) > 0:
         result.fixed = parsed.sql()
-    return { "allowed": len(result.errors) == 0, "errors": result.errors, "fixed": result.fixed, "risk": result.risk}
-
-
-def _verify_where_clause(context: _VerificationContext, select_statement: expr.Query,
-                         from_tables: List[expr.Table]):
+    return {
+        "allowed": len(result.errors) == 0,
+        "errors": result.errors,
+        "fixed": result.fixed,
+        "risk": result.risk,
+    }
+
+
+def _verify_where_clause(
+    context: _VerificationContext,
+    select_statement: expr.Query,
+    from_tables: List[expr.Table],
+):
     _verify_static_expression(select_statement, context)
     _verify_restrictions(select_statement, context, from_tables)
 
-def _verify_restrictions(select_statement: expr.Query,
-                         context: _VerificationContext,
-                         from_tables: List[expr.Table]):
+
+def _verify_restrictions(
+    select_statement: expr.Query,
+    context: _VerificationContext,
+    from_tables: List[expr.Table],
+):
     where_clause = select_statement.find(expr.Where)
     if where_clause is None:
         where_clause = select_statement.find(expr.Where)
         and_exps = []
     else:
         and_exps = list(_split_to_expressions(where_clause.this, expr.And))
-    for t in [c_t for c_t in context.config["tables"] if c_t["table_name"] in [t.name for t in from_tables]]:
+    for t in [
+        c_t
+        for c_t in context.config["tables"]
+        if c_t["table_name"] in [t.name for t in from_tables]
+    ]:
         for idx, r in enumerate(t.get("restrictions", [])):
             found = False
             for sub_exp in and_exps:
@@ -129,18 +146,30 @@ def _verify_restrictions(select_statement: expr.Query,
             if not found:
                 context.add_error(
                     f"Missing restriction for table: {t['table_name']} column: {r['column']} value: {r['value']}",
-                    True, 0.5)
+                    True,
+                    0.5,
+                )
                 value = f"'{r['value']}'" if isinstance(r["value"], str) else r["value"]
-                new_condition = sqlglot.parse_one(f"{r['column']} = {value}", dialect=context.dialect)
+                new_condition = sqlglot.parse_one(
+                    f"{r['column']} = {value}", dialect=context.dialect
+                )
                 if where_clause is None:
                     where_clause = expr.Where(this=new_condition)
                     select_statement.set("where", where_clause)
                 else:
-                    where_clause = where_clause.replace(expr.Where(this=expr.And(this=expr.paren(where_clause.this),
-                                                                                 expression=new_condition)))
-
-
-def _verify_static_expression(select_statement: expr.Query, context: _VerificationContext) -> bool:
+                    where_clause = where_clause.replace(
+                        expr.Where(
+                            this=expr.And(
+                                this=expr.paren(where_clause.this),
+                                expression=new_condition,
+                            )
+                        )
+                    )
+
+
+def _verify_static_expression(
+    select_statement: expr.Query, context: _VerificationContext
+) -> bool:
     has_static_exp = False
     where_clause = select_statement.find(expr.Where)
     if where_clause:
@@ -152,6 +181,7 @@ def _verify_static_expression(select_statement: expr.Query, context: _Verificati
         simplify(where_clause)
     return not has_static_exp
 
+
 def _has_static_expression(context: _VerificationContext, exp: expr.Expression) -> bool:
     if isinstance(exp, expr.Not):
         return _has_static_expression(context, exp.this)
@@ -162,7 +192,8 @@ def _has_static_expression(context: _VerificationContext, exp: expr.Expression)
             result = _has_static_expression(context, sub_exp)
         elif not sub_exp.find(expr.Column):
             context.add_error(
-                f"Static expression is not allowed: {sub_exp.sql()}", True, 0.8)
+                f"Static expression is not allowed: {sub_exp.sql()}", True, 0.8
+            )
             par = sub_exp.parent
             while isinstance(par, expr.Paren):
                 par = par.parent
@@ -176,15 +207,15 @@ def _has_static_expression(context: _VerificationContext, exp: expr.Expression)
 
 def _verify_restriction(restriction: dict, exp: expr.Expression) -> bool:
     """
-       Verifies if a given restriction is satisfied within an SQL expression.
+    Verifies if a given restriction is satisfied within an SQL expression.
 
-       Args:
-           restriction (dict): The restriction to verify, containing 'column' and 'value' keys.
-           exp (list): The SQL expression to check against the restriction.
+    Args:
+        restriction (dict): The restriction to verify, containing 'column' and 'value' keys.
+        exp (list): The SQL expression to check against the restriction.
 
-       Returns:
-           bool: True if the restriction is satisfied, False otherwise.
-   """
+    Returns:
+        bool: True if the restriction is satisfied, False otherwise.
+    """
     if isinstance(exp, expr.Not):
         return False
     if isinstance(exp, expr.Paren):
@@ -202,8 +233,8 @@ def _verify_restriction(restriction: dict, exp: expr.Expression) -> bool:
             return exp.right.this == str(restriction["value"])
     return False
 
-def _verify_query_statement(query_statement: expr.Query,
-                            context: _VerificationContext):
+
+def _verify_query_statement(query_statement: expr.Query, context: _VerificationContext):
     if isinstance(query_statement, expr.Union):
         _verify_query_statement(query_statement.left, context)
         _verify_query_statement(query_statement.right, context)
@@ -226,9 +257,11 @@ def _verify_query_statement(query_statement: expr.Query,
     return query_statement
 
 
-def _verify_select_clause(context: _VerificationContext,
-                          select_clause: expr.Query,
-                          from_tables: List[expr.Table]):
+def _verify_select_clause(
+    context: _VerificationContext,
+    select_clause: expr.Query,
+    from_tables: List[expr.Table],
+):
     to_remove = []
     for e in select_clause.expressions:
         if not _verify_select_clause_element(from_tables, context, e):
@@ -238,8 +271,10 @@ def _verify_select_clause(context: _VerificationContext,
     if len(select_clause.expressions) == 0:
         context.add_error("No legal elements in SELECT clause", False, 0.5)
 
-def _verify_select_clause_element(from_tables: List[expr.Table], context: _VerificationContext,
-                                  e: expr.Expression):
+
+def _verify_select_clause_element(
+    from_tables: List[expr.Table], context: _VerificationContext, e: expr.Expression
+):
     if isinstance(e, expr.Column):
         if not _verify_col(e, from_tables, context):
             return False
@@ -249,7 +284,9 @@ def _verify_select_clause_element(from_tables: List[expr.Table], context: _Verif
             for config_t in context.config["tables"]:
                 if t.name == config_t["table_name"]:
                     for c in config_t["columns"]:
-                        e.parent.set("expressions", e.parent.expressions + [sqlglot.parse_one(c)])
+                        e.parent.set(
+                            "expressions", e.parent.expressions + [sqlglot.parse_one(c)]
+                        )
         return False
     elif isinstance(e, expr.Tuple):
         result = True
@@ -263,7 +300,10 @@ def _verify_select_clause_element(from_tables: List[expr.Table], context: _Verif
                 return False
     return True
 
-def _verify_col(col: expr.Column, from_tables: List[expr.Table], context: _VerificationContext) -> bool:
+
+def _verify_col(
+    col: expr.Column, from_tables: List[expr.Table], context: _VerificationContext
+) -> bool:
     """
     Verifies if a column reference is allowed based on the provided tables and context.
 
@@ -275,16 +315,25 @@ def _verify_col(col: expr.Column, from_tables: List[expr.Table], context: _Verif
     Returns:
         bool: True if the column reference is allowed, False otherwise.
     """
-    if col.table == "sub_select" or  col.table != "" and col.table in context.dynamic_tables:
+    if (
+        col.table == "sub_select"
+        or col.table != ""
+        and col.table in context.dynamic_tables
+    ):
         pass
     elif not _find_column(col.name, from_tables, context):
-        context.add_error(f"Column {col.name} is not allowed. Column removed from SELECT clause",
-                          True,0.3)
+        context.add_error(
+            f"Column {col.name} is not allowed. Column removed from SELECT clause",
+            True,
+            0.3,
+        )
         return False
     return True
 
 
-def _find_column(col_name: str, from_tables: List[expr.Table], result: _VerificationContext) -> bool:
+def _find_column(
+    col_name: str, from_tables: List[expr.Table], result: _VerificationContext
+) -> bool:
     """
     Finds a column in the given tables based on the provided column name.
 
@@ -306,16 +355,18 @@ def _find_column(col_name: str, from_tables: List[expr.Table], result: _Verifica
     return False
 
 
-def _get_from_clause_tables(select_clause: expr.Query, context: _VerificationContext) -> List[expr.Table]:
+def _get_from_clause_tables(
+    select_clause: expr.Query, context: _VerificationContext
+) -> List[expr.Table]:
     """
-        Extracts table references from the FROM clause of an SQL query.
+    Extracts table references from the FROM clause of an SQL query.
 
-        Args:
-            select_clause (dict): The FROM clause of the SQL query.
-            context (_VerificationContext): The context for verification.
+    Args:
+        select_clause (dict): The FROM clause of the SQL query.
+        context (_VerificationContext): The context for verification.
 
-        Returns:
-            List[_TableRef]: A list of table references to find in the FROM clause.
+    Returns:
+        List[_TableRef]: A list of table references to find in the FROM clause.
     """
     result = []
     from_clause = select_clause.find(expr.From)
@@ -337,13 +388,15 @@ def _get_from_clause_tables(select_clause: expr.Query, context: _VerificationCon
     return result
 
 
-def _split_to_expressions(exp: expr.Expression,
-                          exp_type: Type[expr.Expression]) -> Generator[expr.Expression, None, None]:
+def _split_to_expressions(
+    exp: expr.Expression, exp_type: Type[expr.Expression]
+) -> Generator[expr.Expression, None, None]:
     if isinstance(exp, exp_type):
         yield from exp.flatten()
     else:
         yield exp
 
+
 def _find_direct(exp: expr.Expression, exp_type: Type[expr.Expression]):
     for child in exp.args.values():
         if isinstance(child, exp_type):

test/test_rest_api_unit.py

@@ -2,6 +2,7 @@
 
 from sql_data_guard.rest import app
 
+
 class TestRestAppErrors:
     def test_verify_sql_method_not_allowed(self):
         result = app.test_client().get("/verify-sql")
@@ -18,39 +19,49 @@ def test_verify_sql_no_sql(self):
         assert result.json == {"error": "Missing 'sql' in request"}
 
     def test_very_sql_no_config(self):
-        result = app.test_client().post("/verify-sql", json={"sql": "SELECT * FROM my_table"})
+        result = app.test_client().post(
+            "/verify-sql", json={"sql": "SELECT * FROM my_table"}
+        )
         assert result.status_code == 400
         assert result.json == {"error": "Missing 'config' in request"}
 
+
 class TestRestAppVerifySql:
     @pytest.fixture(scope="class")
     def config(self) -> dict:
-        return { "tables": [
-                    {
-                        "table_name": "orders",
-                        "database_name": "orders_db",
-                        "columns": ["id", "product_name", "account_id", "day"],
-                        "restrictions": [{"column": "id", "value": 123}]
-                    }
-                ]
-            }
+        return {
+            "tables": [
+                {
+                    "table_name": "orders",
+                    "database_name": "orders_db",
+                    "columns": ["id", "product_name", "account_id", "day"],
+                    "restrictions": [{"column": "id", "value": 123}],
+                }
+            ]
+        }
 
     def test_verify_sql(self, config):
-        result = app.test_client().post("/verify-sql",
-                                        json={"sql": "SELECT id FROM orders WHERE id = 123",
-                                              "config": config})
+        result = app.test_client().post(
+            "/verify-sql",
+            json={"sql": "SELECT id FROM orders WHERE id = 123", "config": config},
+        )
         assert result.status_code == 200
-        assert result.json == {'allowed': True, 'errors': [], 'fixed': None, "risk": 0}
-
+        assert result.json == {"allowed": True, "errors": [], "fixed": None, "risk": 0}
 
     def test_verify_sql_error(self, config):
-        result = app.test_client().post("/verify-sql",
-                                        json={"sql": "SELECT id, another_col FROM orders WHERE id = 123",
-                                              "config": config})
+        result = app.test_client().post(
+            "/verify-sql",
+            json={
+                "sql": "SELECT id, another_col FROM orders WHERE id = 123",
+                "config": config,
+            },
+        )
         assert result.status_code == 200
         assert result.json == {
-            'allowed': False,
-            'errors': ['Column another_col is not allowed. Column removed from SELECT clause'],
-            'fixed': 'SELECT id FROM orders WHERE id = 123',
-            'risk': 0.3
-       }
+            "allowed": False,
+            "errors": [
+                "Column another_col is not allowed. Column removed from SELECT clause"
+            ],
+            "fixed": "SELECT id FROM orders WHERE id = 123",
+            "risk": 0.3,
+        }