File based column mapping for JDBC connectors

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/mapping/CachingIdentifierMapping.java

@@ -44,6 +44,7 @@ public final class CachingIdentifierMapping
 {
     private final NonKeyEvictableCache<ConnectorIdentity, Mapping> remoteSchemaNames;
     private final NonKeyEvictableCache<RemoteTableNameCacheKey, Mapping> remoteTableNames;
+    private final NonKeyEvictableCache<RemoteColumnNameCacheKey, Mapping> remoteColumnNames;
     private final IdentifierMapping identifierMapping;
 
     @Inject
@@ -55,6 +56,7 @@ public CachingIdentifierMapping(
                 .expireAfterWrite(mappingConfig.getCaseInsensitiveNameMatchingCacheTtl().toMillis(), MILLISECONDS);
         this.remoteSchemaNames = buildNonEvictableCacheWithWeakInvalidateAll(remoteNamesCacheBuilder);
         this.remoteTableNames = buildNonEvictableCacheWithWeakInvalidateAll(remoteNamesCacheBuilder);
+        this.remoteColumnNames = buildNonEvictableCacheWithWeakInvalidateAll(remoteNamesCacheBuilder);
 
         this.identifierMapping = requireNonNull(identifierMapping, "identifierMapping is null");
     }
@@ -80,9 +82,9 @@ public String fromRemoteTableName(String remoteSchemaName, String remoteTableNam
     }
 
     @Override
-    public String fromRemoteColumnName(String remoteColumnName)
+    public String fromRemoteColumnName(String remoteSchemaName, String remoteTableName, String remoteColumnName)
     {
-        return identifierMapping.fromRemoteColumnName(remoteColumnName);
+        return identifierMapping.fromRemoteColumnName(remoteSchemaName, remoteTableName, remoteColumnName);
     }
 
     @Override
@@ -142,9 +144,38 @@ public String toRemoteTableName(RemoteIdentifiers remoteIdentifiers, ConnectorId
     }
 
     @Override
-    public String toRemoteColumnName(RemoteIdentifiers remoteIdentifiers, String columnName)
+    public String toRemoteColumnName(RemoteIdentifiers remoteIdentifiers, ConnectorIdentity identity, String remoteSchema, String remoteTable, String columnName)
     {
-        return identifierMapping.toRemoteColumnName(remoteIdentifiers, columnName);
+        requireNonNull(remoteSchema, "remoteSchema is null");
+        requireNonNull(remoteTable, "remoteTable is null");
+        verify(CharMatcher.forPredicate(Character::isUpperCase).matchesNoneOf(columnName), "Expected column name from internal metadata to be lowercase: %s", columnName);
+        try {
+            RemoteColumnNameCacheKey cacheKey = new RemoteColumnNameCacheKey(identity, remoteSchema, remoteTable);
+            Mapping mapping = remoteColumnNames.getIfPresent(cacheKey);
+            if (mapping != null && !mapping.hasRemoteObject(columnName)) {
+                mapping = null;
+            }
+            if (mapping == null) {
+                mapping = createColumnMapping(remoteSchema, remoteTable, remoteIdentifiers.getRemoteColumns(remoteSchema, remoteTable));
+                remoteColumnNames.put(cacheKey, mapping);
+            }
+            String remoteColumn = mapping.get(columnName);
+            if (remoteColumn != null) {
+                return remoteColumn;
+            }
+        }
+        catch (RuntimeException e) {
+            throw new TrinoException(GENERIC_INTERNAL_ERROR, "Failed to find remote column name: " + firstNonNull(e.getMessage(), e), e);
+        }
+
+        return identifierMapping.toRemoteColumnName(remoteIdentifiers, identity, remoteSchema, remoteTable, columnName);
+    }
+
+    private Mapping createColumnMapping(String remoteSchema, String remoteTable, Set<String> remoteColumns)
+    {
+        return createMapping(
+                remoteColumns,
+                remoteColumnName -> identifierMapping.fromRemoteColumnName(remoteSchema, remoteTable, remoteColumnName));
     }
 
     private Mapping createSchemaMapping(Collection<String> remoteSchemas)

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/mapping/ColumnMappingRule.java

@@ -0,0 +1,88 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.base.mapping;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Objects;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static java.util.Locale.ENGLISH;
+import static java.util.Objects.requireNonNull;
+
+public class ColumnMappingRule
+{
+    private final String remoteSchema;
+    private final String remoteTable;
+    private final String remoteColumn;
+    private final String mapping;
+
+    @JsonCreator
+    public ColumnMappingRule(
+            @JsonProperty String remoteSchema,
+            @JsonProperty String remoteTable,
+            @JsonProperty String remoteColumn,
+            @JsonProperty String mapping)
+    {
+        this.remoteSchema = requireNonNull(remoteSchema, "remoteSchema is null");
+        this.remoteTable = requireNonNull(remoteTable, "remoteTable is null");
+        this.remoteColumn = requireNonNull(remoteColumn, "remoteColumn is null");
+        this.mapping = requireNonNull(mapping, "mapping is null");
+        checkArgument(mapping.toLowerCase(ENGLISH).equals(mapping), "Mapping is not lower cased: %s", mapping);
+    }
+
+    @JsonProperty
+    public String getRemoteSchema()
+    {
+        return remoteSchema;
+    }
+
+    @JsonProperty
+    public String getRemoteTable()
+    {
+        return remoteTable;
+    }
+
+    @JsonProperty
+    public String getRemoteColumn()
+    {
+        return remoteColumn;
+    }
+
+    @JsonProperty
+    public String getMapping()
+    {
+        return mapping;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        ColumnMappingRule that = (ColumnMappingRule) o;
+        return remoteSchema.equals(that.remoteSchema) && remoteTable.equals(that.remoteTable) && remoteColumn.equals(that.remoteColumn) && mapping.equals(that.mapping);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(remoteSchema, remoteTable, mapping);
+    }
+}

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/mapping/DefaultIdentifierMapping.java

@@ -33,7 +33,7 @@ public String fromRemoteTableName(String remoteSchemaName, String remoteTableNam
     }
 
     @Override
-    public String fromRemoteColumnName(String remoteColumnName)
+    public String fromRemoteColumnName(String remoteSchemaName, String remoteTableName, String remoteColumnName)
     {
         return remoteColumnName.toLowerCase(ENGLISH);
     }
@@ -51,7 +51,7 @@ public String toRemoteTableName(RemoteIdentifiers remoteIdentifiers, ConnectorId
     }
 
     @Override
-    public String toRemoteColumnName(RemoteIdentifiers remoteIdentifiers, String columnName)
+    public String toRemoteColumnName(RemoteIdentifiers remoteIdentifiers, ConnectorIdentity identity, String remoteSchemaName, String remoteTableName, String columnName)
     {
         return toRemoteIdentifier(columnName, remoteIdentifiers);
     }

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/mapping/ForwardingIdentifierMapping.java

@@ -51,9 +51,9 @@ public String fromRemoteTableName(String remoteSchemaName, String remoteTableNam
     }
 
     @Override
-    public String fromRemoteColumnName(String remoteColumnName)
+    public String fromRemoteColumnName(String remoteSchemaName, String remoteTableName, String remoteColumnName)
     {
-        return delegate().fromRemoteColumnName(remoteColumnName);
+        return delegate().fromRemoteColumnName(remoteSchemaName, remoteTableName, remoteColumnName);
     }
 
     @Override
@@ -69,8 +69,8 @@ public String toRemoteTableName(RemoteIdentifiers remoteIdentifiers, ConnectorId
     }
 
     @Override
-    public String toRemoteColumnName(RemoteIdentifiers remoteIdentifiers, String columnName)
+    public String toRemoteColumnName(RemoteIdentifiers remoteIdentifiers, ConnectorIdentity identity, String remoteSchemaName, String remoteTableName, String columnName)
     {
-        return delegate().toRemoteColumnName(remoteIdentifiers, columnName);
+        return delegate().toRemoteColumnName(remoteIdentifiers, identity, remoteSchemaName, remoteTableName, columnName);
     }
 }

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/mapping/IdentifierMapping.java

@@ -21,11 +21,11 @@ public interface IdentifierMapping
 
     String fromRemoteTableName(String remoteSchemaName, String remoteTableName);
 
-    String fromRemoteColumnName(String remoteColumnName);
+    String fromRemoteColumnName(String remoteSchemaName, String remoteTableName, String remoteColumnName);
 
     String toRemoteSchemaName(RemoteIdentifiers remoteIdentifiers, ConnectorIdentity identity, String schemaName);
 
     String toRemoteTableName(RemoteIdentifiers remoteIdentifiers, ConnectorIdentity identity, String remoteSchema, String tableName);
 
-    String toRemoteColumnName(RemoteIdentifiers remoteIdentifiers, String columnName);
+    String toRemoteColumnName(RemoteIdentifiers remoteIdentifiers, ConnectorIdentity identity, String remoteSchemaName, String remoteTableName, String columnName);
 }

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/mapping/IdentifierMappingRules.java

@@ -17,22 +17,26 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 import java.util.List;
-import java.util.Objects;
 
+import static java.util.Objects.hash;
 import static java.util.Objects.requireNonNull;
+import static java.util.Objects.requireNonNullElseGet;
 
 public class IdentifierMappingRules
 {
     private final List<SchemaMappingRule> schemas;
     private final List<TableMappingRule> tables;
+    private final List<ColumnMappingRule> columns;
 
     @JsonCreator
     public IdentifierMappingRules(
             @JsonProperty("schemas") List<SchemaMappingRule> schemas,
-            @JsonProperty("tables") List<TableMappingRule> tables)
+            @JsonProperty("tables") List<TableMappingRule> tables,
+            @JsonProperty("columns") List<ColumnMappingRule> columns)
     {
         this.schemas = requireNonNull(schemas, "schemaMappingRules is null");
         this.tables = requireNonNull(tables, "tableMappingRules is null");
+        this.columns = requireNonNullElseGet(columns, List::of);
     }
 
     @JsonProperty("schemas")
@@ -47,6 +51,12 @@ public List<TableMappingRule> getTableMapping()
         return tables;
     }
 
+    @JsonProperty("columns")
+    public List<ColumnMappingRule> getColumns()
+    {
+        return columns;
+    }
+
     @Override
     public boolean equals(Object o)
     {
@@ -57,12 +67,12 @@ public boolean equals(Object o)
             return false;
         }
         IdentifierMappingRules that = (IdentifierMappingRules) o;
-        return schemas.equals(that.schemas) && tables.equals(that.tables);
+        return schemas.equals(that.schemas) && tables.equals(that.tables) && columns.equals(that.columns);
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(schemas, tables);
+        return hash(schemas, tables, columns);
     }
 }

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/mapping/RemoteColumnNameCacheKey.java

@@ -0,0 +1,66 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.base.mapping;
+
+import io.trino.spi.security.ConnectorIdentity;
+
+import java.util.Objects;
+
+import static com.google.common.base.MoreObjects.toStringHelper;
+import static java.util.Objects.requireNonNull;
+
+public class RemoteColumnNameCacheKey
+{
+    private final ConnectorIdentity identity;
+    private final String schema;
+    private final String table;
+
+    RemoteColumnNameCacheKey(ConnectorIdentity identity, String schema, String table)
+    {
+        this.identity = requireNonNull(identity, "identity is null");
+        this.schema = requireNonNull(schema, "schema is null");
+        this.table = requireNonNull(table, "table is null");
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        RemoteColumnNameCacheKey that = (RemoteColumnNameCacheKey) o;
+        return Objects.equals(identity, that.identity) &&
+                Objects.equals(schema, that.schema) &&
+                Objects.equals(table, that.table);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(identity, schema, table);
+    }
+
+    @Override
+    public String toString()
+    {
+        return toStringHelper(this)
+                .add("identity", identity)
+                .add("schema", schema)
+                .add("table", table)
+                .toString();
+    }
+}

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/mapping/RemoteIdentifiers.java

@@ -21,5 +21,7 @@ public interface RemoteIdentifiers
 
     Set<String> getRemoteTables(String remoteSchema);
 
+    Set<String> getRemoteColumns(String remoteSchema, String remoteTable);
+
     boolean storesUpperCaseIdentifiers();
 }