singhpk234 commented on code in PR #1802:
URL: https://github.com/apache/polaris/pull/1802#discussion_r2129038100


##########
persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/QueryGenerator.java:
##########
@@ -21,193 +21,209 @@
 import com.google.common.annotations.VisibleForTesting;
 import jakarta.annotation.Nonnull;
 import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import org.apache.polaris.core.entity.PolarisBaseEntity;
+import java.util.stream.Collectors;
 import org.apache.polaris.core.entity.PolarisEntityCore;
 import org.apache.polaris.core.entity.PolarisEntityId;
-import org.apache.polaris.core.entity.PolarisEntityType;
-import org.apache.polaris.core.policy.PolicyEntity;
-import org.apache.polaris.persistence.relational.jdbc.models.Converter;
 import org.apache.polaris.persistence.relational.jdbc.models.ModelEntity;
 import org.apache.polaris.persistence.relational.jdbc.models.ModelGrantRecord;
-import 
org.apache.polaris.persistence.relational.jdbc.models.ModelPolicyMappingRecord;
-import 
org.apache.polaris.persistence.relational.jdbc.models.ModelPrincipalAuthenticationData;
 
+/**
+ * Utility class to generate parameterized SQL queries (SELECT, INSERT, 
UPDATE, DELETE). Ensures
+ * consistent SQL generation and protects against injection by managing 
parameters separately.
+ */
 public class QueryGenerator {
+  private final DatabaseType databaseType;
 
-  public static <T> String generateSelectQuery(
-      @Nonnull Converter<T> entity, @Nonnull Map<String, Object> whereClause) {
-    return generateSelectQuery(entity, generateWhereClause(whereClause));
+  public QueryGenerator(DatabaseType databaseType) {
+    this.databaseType = databaseType;
   }
 
-  public static String generateDeleteQueryForEntityGrantRecords(
-      @Nonnull PolarisEntityCore entity, @Nonnull String realmId) {
-    String granteeCondition =
-        String.format(
-            "grantee_id = %s AND grantee_catalog_id = %s", entity.getId(), 
entity.getCatalogId());
-    String securableCondition =
-        String.format(
-            "securable_id = %s AND securable_catalog_id = %s",
-            entity.getId(), entity.getCatalogId());
-
-    String whereClause =
-        " WHERE ("
-            + granteeCondition
-            + " OR "
-            + securableCondition
-            + ") AND realm_id = '"
-            + realmId
-            + "'";
-    return generateDeleteQuery(ModelGrantRecord.class, whereClause);
+  /**
+   * @return The configured database type.
+   */
+  public DatabaseType getDatabaseType() {
+    return databaseType;
   }
 
-  public static String generateDeleteQueryForEntityPolicyMappingRecords(
-      @Nonnull PolarisBaseEntity entity, @Nonnull String realmId) {
-    Map<String, Object> queryParams = new HashMap<>();
-    if (entity.getType() == PolarisEntityType.POLICY) {
-      PolicyEntity policyEntity = PolicyEntity.of(entity);
-      queryParams.put("policy_type_code", policyEntity.getPolicyTypeCode());
-      queryParams.put("policy_catalog_id", policyEntity.getCatalogId());
-      queryParams.put("policy_id", policyEntity.getId());
-    } else {
-      queryParams.put("target_catalog_id", entity.getCatalogId());
-      queryParams.put("target_id", entity.getId());
-    }
-    queryParams.put("realm_id", realmId);
+  /** A container for the SQL string and the ordered parameter values. */
+  public record PreparedQuery(String sql, List<Object> parameters) {}
+
+  /**
+   * Generates a SELECT query with projection and filtering.
+   *
+   * @param projections List of columns to retrieve.
+   * @param tableName Target table name.
+   * @param whereClause Column-value pairs used in WHERE filtering.
+   * @return A parameterized SELECT query.
+   * @throws IllegalArgumentException if any whereClause column isn't in 
projections.
+   */
+  public PreparedQuery generateSelectQuery(
+      @Nonnull List<String> projections,
+      @Nonnull String tableName,
+      @Nonnull Map<String, Object> whereClause) {
+    checkInvalidColumns(projections, whereClause);
+    PreparedQuery where = generateWhereClause(whereClause);
+    PreparedQuery query = generateSelectQuery(projections, tableName, 
where.sql());
+    return new PreparedQuery(query.sql(), where.parameters());
+  }
 
-    return generateDeleteQuery(ModelPolicyMappingRecord.class, queryParams);
+  /**
+   * Builds a DELETE query to remove grant records for a given entity.
+   *
+   * @param entity The target entity (either grantee or securable).
+   * @param realmId The associated realm.
+   * @return A DELETE query removing all grants for this entity.
+   */
+  public PreparedQuery generateDeleteQueryForEntityGrantRecords(
+      @Nonnull PolarisEntityCore entity, @Nonnull String realmId) {
+    String where =
+        " WHERE ((grantee_id = ? AND grantee_catalog_id = ?) OR (securable_id 
= ? AND securable_catalog_id = ?)) AND realm_id = ?";
+    List<Object> params =
+        Arrays.asList(
+            entity.getId(), entity.getCatalogId(), entity.getId(), 
entity.getCatalogId(), realmId);
+    return new PreparedQuery(
+        "DELETE FROM " + 
getFullyQualifiedTableName(ModelGrantRecord.TABLE_NAME) + where, params);
   }
 
-  public static String generateSelectQueryWithEntityIds(
+  /**
+   * Builds a SELECT query using a list of entity ID pairs (catalog_id, id).
+   *
+   * @param realmId Realm to filter by.
+   * @param entityIds List of PolarisEntityId pairs.
+   * @return SELECT query to retrieve matching entities.
+   * @throws IllegalArgumentException if entityIds is empty.
+   */
+  public PreparedQuery generateSelectQueryWithEntityIds(
       @Nonnull String realmId, @Nonnull List<PolarisEntityId> entityIds) {
     if (entityIds.isEmpty()) {
       throw new IllegalArgumentException("Empty entity ids");
     }
-    StringBuilder condition = new StringBuilder("(catalog_id, id) IN (");
-    for (PolarisEntityId entityId : entityIds) {
-      String in = "(" + entityId.getCatalogId() + ", " + entityId.getId() + 
")";
-      condition.append(in);
-      condition.append(",");
+    String placeholders = entityIds.stream().map(e -> "(?, 
?)").collect(Collectors.joining(", "));
+    List<Object> params = new ArrayList<>();
+    for (PolarisEntityId id : entityIds) {
+      params.add(id.getCatalogId());
+      params.add(id.getId());
     }
-    // extra , removed
-    condition.deleteCharAt(condition.length() - 1);
-    condition.append(")");
-    condition.append(" AND realm_id = '").append(realmId).append("'");
-
-    return generateSelectQuery(new ModelEntity(), " WHERE " + condition);
+    params.add(realmId);
+    String where = " WHERE (catalog_id, id) IN (" + placeholders + ") AND 
realm_id = ?";
+    return new PreparedQuery(
+        generateSelectQuery(ModelEntity.ALL_COLUMNS, ModelEntity.TABLE_NAME, 
where).sql(), params);
   }
 
-  public static <T> String generateInsertQuery(
-      @Nonnull Converter<T> entity, @Nonnull String realmId) {
-    String tableName = getTableName(entity.getClass());
-    Map<String, Object> obj = entity.toMap();
-    List<String> columnNames = new ArrayList<>(obj.keySet());
-    List<String> values =
-        new ArrayList<>(obj.values().stream().map(val -> "'" + val.toString() 
+ "'").toList());
-    columnNames.add("realm_id");
-    values.add("'" + realmId + "'");
-
-    String columns = String.join(", ", columnNames);
-    String valuesString = String.join(", ", values);
-
-    return "INSERT INTO " + tableName + " (" + columns + ") VALUES (" + 
valuesString + ")";
+  /**
+   * Generates an INSERT query for a given table.
+   *
+   * @param allColumns Columns to insert values into.
+   * @param tableName Target table name.
+   * @param values Values for each column (must match order of columns).
+   * @param realmId Realm value to append.
+   * @return INSERT query with value bindings.
+   */
+  public <T> PreparedQuery generateInsertQuery(
+      @Nonnull List<String> allColumns,
+      @Nonnull String tableName,
+      List<Object> values,
+      String realmId) {
+    List<String> finalColumns = new ArrayList<>(allColumns);
+    List<Object> finalValues = new ArrayList<>(values);
+    finalColumns.add("realm_id");
+    finalValues.add(realmId);
+    String columns = String.join(", ", finalColumns);
+    String placeholders = finalColumns.stream().map(c -> 
"?").collect(Collectors.joining(", "));
+    String sql =
+        "INSERT INTO "
+            + getFullyQualifiedTableName(tableName)
+            + " ("
+            + columns
+            + ") VALUES ("
+            + placeholders
+            + ")";
+    return new PreparedQuery(sql, finalValues);
   }
 
-  public static <T> String generateUpdateQuery(
-      @Nonnull Converter<T> entity, @Nonnull Map<String, Object> whereClause) {
-    String tableName = getTableName(entity.getClass());
-    Map<String, Object> obj = entity.toMap();
-    List<String> setClauses = new ArrayList<>();
-    List<String> columnNames = new ArrayList<>(obj.keySet());
-    List<String> values = obj.values().stream().map(val -> "'" + 
val.toString() + "'").toList();
-
-    for (int i = 0; i < columnNames.size(); i++) {
-      setClauses.add(columnNames.get(i) + " = " + values.get(i)); // 
Placeholders
-    }
-
-    String setClausesString = String.join(", ", setClauses);
-
-    return "UPDATE " + tableName + " SET " + setClausesString + 
generateWhereClause(whereClause);
+  /**
+   * Builds an UPDATE query.
+   *
+   * @param allColumns Columns to update.
+   * @param tableName Target table.
+   * @param values New values (must match columns in order).
+   * @param whereClause Conditions for filtering rows to update.
+   * @return UPDATE query with parameter values.
+   */
+  public <T> PreparedQuery generateUpdateQuery(
+      @Nonnull List<String> allColumns,
+      @Nonnull String tableName,
+      @Nonnull List<Object> values,
+      @Nonnull Map<String, Object> whereClause) {
+    checkInvalidColumns(allColumns, whereClause);
+    List<Object> bindingParams = new ArrayList<>(values);
+    PreparedQuery where = generateWhereClause(whereClause);
+    String setClause = allColumns.stream().map(c -> c + " = 
?").collect(Collectors.joining(", "));
+    String sql =
+        "UPDATE " + getFullyQualifiedTableName(tableName) + " SET " + 
setClause + where.sql();
+    bindingParams.addAll(where.parameters());
+    return new PreparedQuery(sql, bindingParams);
   }
 
-  public static String generateDeleteQuery(
-      @Nonnull Class<?> entityClass, @Nonnull Map<String, Object> whereClause) 
{
-    return generateDeleteQuery(entityClass, 
(generateWhereClause(whereClause)));
+  /**
+   * Builds a DELETE query with the given conditions.
+   *
+   * @param tableColumns List of valid table columns.
+   * @param tableName Target table.
+   * @param whereClause Column-value filters.
+   * @return DELETE query with parameter bindings.
+   */
+  public PreparedQuery generateDeleteQuery(
+      @Nonnull List<String> tableColumns,
+      @Nonnull String tableName,
+      @Nonnull Map<String, Object> whereClause) {
+    checkInvalidColumns(tableColumns, whereClause);
+    PreparedQuery where = generateWhereClause(whereClause);
+    return new PreparedQuery(
+        "DELETE FROM " + getFullyQualifiedTableName(tableName) + where.sql(), 
where.parameters());
   }
 
-  public static String generateDeleteQuery(
-      @Nonnull Class<?> entityClass, @Nonnull String whereClause) {
-    return "DELETE FROM " + getTableName(entityClass) + whereClause;
-  }
-
-  public static String generateDeleteAll(@Nonnull Class<?> entityClass, 
@Nonnull String realmId) {
-    String tableName = getTableName(entityClass);
-    return "DELETE FROM " + tableName + " WHERE 1 = 1 AND realm_id = '" + 
realmId + "'";
-  }
-
-  public static <T> String generateDeleteQuery(
-      @Nonnull Converter<T> entity, @Nonnull String realmId) {
-    String tableName = getTableName(entity.getClass());
-    Map<String, Object> objMap = entity.toMap();
-    objMap.put("realm_id", realmId);
-    String whereConditions = generateWhereClause(objMap);
-    return "DELETE FROM " + tableName + whereConditions;
+  @VisibleForTesting
+  PreparedQuery generateSelectQuery(
+      @Nonnull List<String> columnNames, @Nonnull String tableName, @Nonnull 
String filter) {

Review Comment:
   This method should have been private and is used strictly within 
QueryGenerator itself, this is made package protected with explicit 
@VisbleforTesting to test coverage, if this is causing some confusion or 
un-ease I am happy to remove this method or test for this method 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to