Skip to content

Commit 7b05f4b

Browse files
patelraj0602patelraj0602
committed
added registry
1 parent df3a484 commit 7b05f4b

File tree

4 files changed

+385
-0
lines changed

4 files changed

+385
-0
lines changed

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnInfo.java

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
package org.hypertrace.core.documentstore.postgres.registry;
2+
3+
import java.util.Objects;
4+
5+
/**
6+
* Immutable data class representing metadata about a PostgreSQL column. Contains the column name
7+
* and its corresponding PostgreSQL data type.
8+
*/
9+
public final class PostgresColumnInfo {
10+
11+
private final String columnName;
12+
private final PostgresColumnType columnType;
13+
14+
/**
15+
* Creates a new PostgresColumnInfo instance.
16+
*
17+
* @param columnName the PostgreSQL column name (must not be null)
18+
* @param columnType the PostgreSQL column type (must not be null)
19+
* @throws IllegalArgumentException if columnName or columnType is null
20+
*/
21+
public PostgresColumnInfo(String columnName, PostgresColumnType columnType) {
22+
this.columnName = Objects.requireNonNull(columnName, "Column name cannot be null");
23+
this.columnType = Objects.requireNonNull(columnType, "Column type cannot be null");
24+
}
25+
26+
/**
27+
* Gets the PostgreSQL column name.
28+
*
29+
* @return the column name
30+
*/
31+
public String getColumnName() {
32+
return columnName;
33+
}
34+
35+
/**
36+
* Gets the PostgreSQL column type.
37+
*
38+
* @return the column type
39+
*/
40+
public PostgresColumnType getColumnType() {
41+
return columnType;
42+
}
43+
44+
/**
45+
* Checks if this column represents a first-class field (non-JSONB).
46+
*
47+
* @return true if this is a first-class field, false if it's a JSONB field
48+
*/
49+
public boolean isFirstClassField() {
50+
return columnType.isFirstClassField();
51+
}
52+
53+
@Override
54+
public boolean equals(Object obj) {
55+
if (this == obj) {
56+
return true;
57+
}
58+
if (obj == null || getClass() != obj.getClass()) {
59+
return false;
60+
}
61+
PostgresColumnInfo that = (PostgresColumnInfo) obj;
62+
return Objects.equals(columnName, that.columnName) && columnType == that.columnType;
63+
}
64+
65+
@Override
66+
public int hashCode() {
67+
return Objects.hash(columnName, columnType);
68+
}
69+
70+
@Override
71+
public String toString() {
72+
return String.format(
73+
"PostgresColumnInfo{columnName='%s', columnType=%s}", columnName, columnType);
74+
}
75+
}

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnRegistry.java

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
package org.hypertrace.core.documentstore.postgres.registry;
2+
3+
import java.util.Optional;
4+
import java.util.Set;
5+
6+
/**
7+
* Registry interface for PostgreSQL column type information. Provides metadata about table columns
8+
* to determine appropriate query generation strategies.
9+
*
10+
* <p>This registry replaces hardcoded column lists and enables dynamic, database-driven decisions
11+
* about whether to use JSON-based or native PostgreSQL column access.
12+
*/
13+
public interface PostgresColumnRegistry {
14+
15+
/**
16+
* Determines if the specified field name corresponds to a first-class column (i.e., a native
17+
* PostgreSQL column rather than a field within a JSONB document).
18+
*
19+
* @param fieldName the field name to check
20+
* @return true if this is a first-class column, false if it should be accessed via JSONB
21+
*/
22+
boolean isFirstClassColumn(String fieldName);
23+
24+
/**
25+
* Gets the PostgreSQL column type for the specified field name.
26+
*
27+
* @param fieldName the field name to look up
28+
* @return the PostgreSQL column type, or empty if the field is not a first-class column
29+
*/
30+
Optional<PostgresColumnType> getColumnType(String fieldName);
31+
32+
/**
33+
* Gets the actual PostgreSQL column name for the specified field name. In most cases, this will
34+
* be the same as the field name, but this method allows for potential field name transformations
35+
* or aliases.
36+
*
37+
* @param fieldName the field name to look up
38+
* @return the PostgreSQL column name, or empty if the field is not a first-class column
39+
*/
40+
Optional<String> getColumnName(String fieldName);
41+
42+
/**
43+
* Gets all first-class column names in this table. This is useful for debugging and validation
44+
* purposes.
45+
*
46+
* @return a set of all first-class column names
47+
*/
48+
Set<String> getAllFirstClassColumns();
49+
50+
/**
51+
* Gets the table name this registry represents.
52+
*
53+
* @return the table name
54+
*/
55+
String getTableName();
56+
57+
/**
58+
* Checks if this registry has any first-class columns. If false, all fields should be accessed
59+
* via JSONB.
60+
*
61+
* @return true if there are first-class columns, false otherwise
62+
*/
63+
default boolean hasFirstClassColumns() {
64+
return !getAllFirstClassColumns().isEmpty();
65+
}
66+
}

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnRegistryImpl.java

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
package org.hypertrace.core.documentstore.postgres.registry;
2+
3+
import java.sql.Connection;
4+
import java.sql.PreparedStatement;
5+
import java.sql.ResultSet;
6+
import java.sql.SQLException;
7+
import java.util.Collections;
8+
import java.util.HashMap;
9+
import java.util.Map;
10+
import java.util.Optional;
11+
import java.util.Set;
12+
import lombok.extern.slf4j.Slf4j;
13+
import org.hypertrace.core.documentstore.postgres.PostgresTableIdentifier;
14+
15+
/**
16+
* Implementation of PostgresColumnRegistry that queries the database schema to build column type
17+
* mappings dynamically.
18+
*
19+
* <p>This implementation queries the information_schema.columns table to discover the actual column
20+
* types in the PostgreSQL table and maps them to the appropriate PostgresColumnType enum values.
21+
*/
22+
@Slf4j
23+
public class PostgresColumnRegistryImpl implements PostgresColumnRegistry {
24+
25+
private final String tableName;
26+
private final Map<String, PostgresColumnInfo> columnMappings;
27+
28+
/**
29+
* Creates a new PostgresColumnRegistryImpl by querying the database schema.
30+
*
31+
* @param connection the database connection to use for schema queries
32+
* @param tableIdentifier the table identifier to query schema for
33+
* @throws SQLException if there's an error querying the database schema
34+
*/
35+
public PostgresColumnRegistryImpl(Connection connection, PostgresTableIdentifier tableIdentifier)
36+
throws SQLException {
37+
this.tableName = tableIdentifier.getTableName();
38+
this.columnMappings = buildColumnMappings(connection, tableIdentifier);
39+
System.out.println("-- PostgresColumnRegistryImpl --");
40+
System.out.printf("Table Name: %s\n", this.tableName);
41+
System.out.println(this.columnMappings);
42+
System.out.println("-- end --");
43+
System.out.println();
44+
45+
log.debug(
46+
"Created PostgresColumnRegistry for table '{}' with {} first-class columns: {}",
47+
tableName,
48+
columnMappings.size(),
49+
columnMappings.keySet());
50+
}
51+
52+
/**
53+
* Creates a new PostgresColumnRegistryImpl by querying the database schema.
54+
*
55+
* @param connection the database connection to use for schema queries
56+
* @param collectionName the collection name (table name) to query schema for
57+
* @throws SQLException if there's an error querying the database schema
58+
*/
59+
public PostgresColumnRegistryImpl(Connection connection, String collectionName)
60+
throws SQLException {
61+
this(connection, PostgresTableIdentifier.parse(collectionName));
62+
}
63+
64+
@Override
65+
public boolean isFirstClassColumn(String fieldName) {
66+
return columnMappings.containsKey(fieldName);
67+
}
68+
69+
@Override
70+
public Optional<PostgresColumnType> getColumnType(String fieldName) {
71+
PostgresColumnInfo columnInfo = columnMappings.get(fieldName);
72+
return columnInfo != null ? Optional.of(columnInfo.getColumnType()) : Optional.empty();
73+
}
74+
75+
@Override
76+
public Optional<String> getColumnName(String fieldName) {
77+
PostgresColumnInfo columnInfo = columnMappings.get(fieldName);
78+
return columnInfo != null ? Optional.of(columnInfo.getColumnName()) : Optional.empty();
79+
}
80+
81+
@Override
82+
public Set<String> getAllFirstClassColumns() {
83+
return Collections.unmodifiableSet(columnMappings.keySet());
84+
}
85+
86+
@Override
87+
public String getTableName() {
88+
return tableName;
89+
}
90+
91+
/**
92+
* Builds the column mappings by querying the database schema.
93+
*
94+
* @param connection the database connection
95+
* @param tableIdentifier the table identifier
96+
* @return a map of field names to PostgresColumnInfo objects
97+
* @throws SQLException if there's an error querying the database
98+
*/
99+
private Map<String, PostgresColumnInfo> buildColumnMappings(
100+
Connection connection, PostgresTableIdentifier tableIdentifier) throws SQLException {
101+
102+
Map<String, PostgresColumnInfo> mappings = new HashMap<>();
103+
104+
String query =
105+
"SELECT column_name, data_type, udt_name "
106+
+ "FROM information_schema.columns "
107+
+ "WHERE table_name = ? AND table_schema = ? "
108+
+ "ORDER BY ordinal_position";
109+
110+
String schemaName = tableIdentifier.getSchema().orElse("public");
111+
String tableName = tableIdentifier.getTableName();
112+
113+
try (PreparedStatement stmt = connection.prepareStatement(query)) {
114+
stmt.setString(1, tableName);
115+
stmt.setString(2, schemaName);
116+
117+
log.debug("Querying schema for table: {}.{}", schemaName, tableName);
118+
119+
try (ResultSet rs = stmt.executeQuery()) {
120+
if (rs != null) {
121+
while (rs.next()) {
122+
String columnName = rs.getString("column_name");
123+
String dataType = rs.getString("data_type");
124+
String udtName = rs.getString("udt_name");
125+
126+
PostgresColumnType columnType = PostgresColumnType.fromPostgresType(dataType, udtName);
127+
128+
// Only include first-class columns (non-JSONB) in the registry
129+
if (columnType.isFirstClassField()) {
130+
PostgresColumnInfo columnInfo = new PostgresColumnInfo(columnName, columnType);
131+
mappings.put(columnName, columnInfo);
132+
133+
log.debug(
134+
"Registered first-class column: {} -> {} ({})", columnName, columnType, dataType);
135+
} else {
136+
log.debug("Skipping JSONB column: {} ({})", columnName, dataType);
137+
}
138+
}
139+
}
140+
}
141+
} catch (SQLException e) {
142+
log.error(
143+
"Failed to query schema for table {}.{}: {}", schemaName, tableName, e.getMessage());
144+
throw e;
145+
}
146+
147+
return mappings;
148+
}
149+
150+
@Override
151+
public String toString() {
152+
return String.format(
153+
"PostgresColumnRegistryImpl{tableName='%s', firstClassColumns=%d}",
154+
tableName, columnMappings.size());
155+
}
156+
}

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/registry/PostgresColumnType.java

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
package org.hypertrace.core.documentstore.postgres.registry;
2+
3+
/**
4+
* Enumeration of PostgreSQL column types supported by the document store. This enum maps PostgreSQL
5+
* data types to their corresponding Java representations and determines the appropriate query
6+
* generation strategy.
7+
*/
8+
public enum PostgresColumnType {
9+
/** PostgreSQL text/varchar types - mapped to Java String */
10+
TEXT,
11+
12+
/** PostgreSQL bigint/int8 types - mapped to Java Long */
13+
BIGINT,
14+
15+
/** PostgreSQL double precision/float8 types - mapped to Java Double */
16+
DOUBLE_PRECISION,
17+
18+
/** PostgreSQL boolean/bool types - mapped to Java Boolean */
19+
BOOLEAN,
20+
21+
/** PostgreSQL text array types - mapped to Java String[] */
22+
TEXT_ARRAY,
23+
24+
/** PostgreSQL jsonb/json types - existing document storage format */
25+
JSONB;
26+
27+
/**
28+
* Determines if this column type represents a first-class (non-JSON) field. First-class fields
29+
* are stored as native PostgreSQL types rather than within JSONB documents.
30+
*
31+
* @return true if this is a first-class field type, false if it's a JSON document field
32+
*/
33+
public boolean isFirstClassField() {
34+
return this != JSONB;
35+
}
36+
37+
/**
38+
* Maps PostgreSQL data type names to PostgresColumnType enum values.
39+
*
40+
* @param dataType the PostgreSQL data type name (e.g., "text", "bigint", "boolean")
41+
* @param udtName the user-defined type name for arrays (e.g., "_text" for text arrays)
42+
* @return the corresponding PostgresColumnType, or JSONB as fallback
43+
*/
44+
public static PostgresColumnType fromPostgresType(String dataType, String udtName) {
45+
if (dataType == null) {
46+
return JSONB;
47+
}
48+
49+
String lowerDataType = dataType.toLowerCase();
50+
51+
// Handle array types first (indicated by udtName starting with underscore)
52+
if (udtName != null && udtName.startsWith("_")) {
53+
switch (udtName.toLowerCase()) {
54+
case "_text":
55+
return TEXT_ARRAY;
56+
default:
57+
return JSONB; // Unsupported array type, fallback to JSONB
58+
}
59+
}
60+
61+
// Handle scalar types
62+
switch (lowerDataType) {
63+
case "text":
64+
case "varchar":
65+
case "character varying":
66+
return TEXT;
67+
68+
case "bigint":
69+
case "int8":
70+
return BIGINT;
71+
72+
case "double precision":
73+
case "float8":
74+
return DOUBLE_PRECISION;
75+
76+
case "boolean":
77+
case "bool":
78+
return BOOLEAN;
79+
80+
case "jsonb":
81+
case "json":
82+
return JSONB;
83+
84+
default:
85+
return JSONB; // Unsupported type, fallback to existing JSONB behavior
86+
}
87+
}
88+
}

0 commit comments

Comments
 (0)