Merge branch 'postgresql-dialect' into allow-t-in-timestamp

src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java

@@ -775,6 +775,26 @@ public String getServerVersion() {
     return serverVersion;
   }
 
+  public String getServerVersionNum() {
+    String[] components = serverVersion.split("\\.");
+    if (components.length >= 2) {
+      int major = tryParseInt(components[0]);
+      int minor = tryParseInt(components[1]);
+      if (major > -1 && minor > -1) {
+        return String.valueOf(major * 10000 + minor);
+      }
+    }
+    return serverVersion;
+  }
+
+  private static int tryParseInt(String value) {
+    try {
+      return Integer.parseInt(value);
+    } catch (NumberFormatException ignore) {
+      return -1;
+    }
+  }
+
   /** Returns true if the OS is Windows. */
   public boolean isWindows() {
     return osName.toLowerCase().startsWith("windows");

src/main/java/com/google/cloud/spanner/pgadapter/session/PGSetting.java

@@ -25,6 +25,7 @@
 import java.util.Arrays;
 import java.util.Objects;
 import java.util.Scanner;
+import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 
 /** Represents a row in the pg_settings table. */
@@ -47,6 +48,26 @@ public class PGSetting {
   private static final int SOURCEFILE_INDEX = 14;
   private static final int SOURCELINE_INDEX = 15;
   private static final int PENDING_RESTART_INDEX = 16;
+  /** The column names of the pg_settings table. */
+  private static final ImmutableList<String> COLUMN_NAMES =
+      ImmutableList.of(
+          "name",
+          "setting",
+          "unit",
+          "category",
+          "short_desc",
+          "extra_desc",
+          "context",
+          "vartype",
+          "source",
+          "min_val",
+          "max_val",
+          "enumvals",
+          "boot_val",
+          "reset_val",
+          "sourcefile",
+          "sourceline",
+          "pending_restart");
 
   private final String extension;
   private final String name;
@@ -178,6 +199,7 @@ private PGSetting(
     this.pendingRestart = pendingRestart;
   }
 
+  /** Returns a copy of this {@link PGSetting}. */
   PGSetting copy() {
     return new PGSetting(
         extension,
@@ -200,7 +222,86 @@ PGSetting copy() {
         pendingRestart);
   }
 
-  public String getKey() {
+  /** Returns this setting as a SELECT statement that can be used in a query or CTE. */
+  String getSelectStatement() {
+    return "select "
+        + toSelectExpression(getCasePreservingKey())
+        + " as name, "
+        + toSelectExpression(setting)
+        + " as setting, "
+        + toSelectExpression(unit)
+        + " as unit, "
+        + toSelectExpression(category)
+        + " as category, "
+        + toSelectExpression((String) null)
+        + " as short_desc, "
+        + toSelectExpression((String) null)
+        + " as extra_desc, "
+        + toSelectExpression(context)
+        + " as context, "
+        + toSelectExpression(vartype)
+        + " as vartype, "
+        + toSelectExpression(minVal)
+        + " as min_val, "
+        + toSelectExpression(maxVal)
+        + " as max_val, "
+        + toSelectExpression(enumVals)
+        + " as enumvals, "
+        + toSelectExpression(bootVal)
+        + " as boot_val, "
+        + toSelectExpression(resetVal)
+        + " as reset_val, "
+        + toSelectExpression(source)
+        + " as source, "
+        + toSelectExpression((String) null)
+        + " as sourcefile, "
+        + toSelectExpression((Integer) null)
+        + "::bigint as sourceline, "
+        + toSelectExpression(pendingRestart)
+        + "::boolean as pending_restart";
+  }
+
+  /** Returns the column names of the pg_settings table. */
+  static ImmutableList<String> getColumnNames() {
+    return COLUMN_NAMES;
+  }
+
+  /** Converts a string to a SQL literal expression that can be used in a select statement. */
+  String toSelectExpression(String value) {
+    return value == null ? "null" : "'" + value + "'";
+  }
+
+  /**
+   * Converts a string array to a SQL literal expression that can be used in a select statement. The
+   * expression is cast to text[].
+   */
+  String toSelectExpression(String[] value) {
+    return value == null
+        ? "null::text[]"
+        : "'{"
+            + Arrays.stream(value)
+                .map(s -> s.startsWith("\"") ? s : "\"" + s + "\"")
+                .collect(Collectors.joining(", "))
+            + "}'::text[]";
+  }
+
+  /** Converts an Integer to a SQL literal expression that can be used in a select statement. */
+  String toSelectExpression(Integer value) {
+    return value == null ? "null" : value.toString();
+  }
+
+  /** Converts a Boolean to a SQL literal expression that can be used in a select statement. */
+  String toSelectExpression(Boolean value) {
+    return value == null ? "null" : (value ? "'t'" : "'f'");
+  }
+
+  /**
+   * Returns the case-preserving key of this setting. Some settings have a key that is written in
+   * camel case (e.g. 'DateStyle') instead of snake case (e.g. 'server_version'). This key should
+   * not be used to look up a setting in the session state map, but should be used in for example
+   * the pg_settings table.
+   */
+  public String getCasePreservingKey() {
     if (extension == null) {
       return name;
     }
@@ -225,6 +326,12 @@ void initSettingValue(String value) {
     this.setting = value;
   }
 
+  /** Initializes the value of the setting at connection startup. */
+  void initConnectionValue(String value) {
+    setSetting(value);
+    this.resetVal = value;
+  }
+
   /**
    * Sets the value for this setting. Throws {@link SpannerException} if the value is not valid, or
    * if the setting is not settable.
@@ -247,7 +354,7 @@ boolean isSettable() {
 
   private void checkValidContext() {
     if (!isSettable()) {
-      throw invalidContextError(getKey(), this.context);
+      throw invalidContextError(getCasePreservingKey(), this.context);
     }
   }
 
@@ -277,22 +384,22 @@ private void checkValidValue(String value) {
       try {
         BooleanParser.toBoolean(value);
       } catch (IllegalArgumentException exception) {
-        throw invalidBoolError(getKey());
+        throw invalidBoolError(getCasePreservingKey());
       }
     } else if ("integer".equals(this.vartype)) {
       try {
         Integer.parseInt(value);
       } catch (NumberFormatException exception) {
-        throw invalidValueError(getKey(), value);
+        throw invalidValueError(getCasePreservingKey(), value);
       }
     } else if ("real".equals(this.vartype)) {
       try {
         Double.parseDouble(value);
       } catch (NumberFormatException exception) {
-        throw invalidValueError(getKey(), value);
+        throw invalidValueError(getCasePreservingKey(), value);
       }
     } else if (enumVals != null && !Iterables.contains(Arrays.asList(this.enumVals), value)) {
-      throw invalidValueError(getKey(), value);
+      throw invalidValueError(getCasePreservingKey(), value);
     }
   }
 

src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java

@@ -18,7 +18,11 @@
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerExceptionFactory;
+import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.Value;
+import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -27,12 +31,37 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 /** {@link SessionState} contains all session variables for a connection. */
 @InternalApi
 public class SessionState {
+  /**
+   * This set contains the settings that show up in the pg_settings CTE. Not all settings are
+   * included in the CTE because Cloud Spanner has a limit of max 60 union all clauses in a
+   * sub-select.
+   */
+  private static final ImmutableSet<String> SUPPORTED_PG_SETTINGS_KEYS =
+      ImmutableSet.of(
+          "application_name",
+          "bytea_output",
+          "DateStyle",
+          "default_transaction_isolation",
+          "default_transaction_read_only",
+          "extra_float_digits",
+          "max_connections",
+          "max_index_keys",
+          "port",
+          "search_path",
+          "server_version",
+          "server_version_num",
+          "TimeZone",
+          "transaction_isolation",
+          "transaction_read_only");
+
   private static final Map<String, PGSetting> SERVER_SETTINGS = new HashMap<>();
 
   static {
@@ -49,8 +78,96 @@ public class SessionState {
   private Map<String, PGSetting> localSettings;
 
   public SessionState(OptionsMetadata options) {
-    this.settings = new HashMap<>(SERVER_SETTINGS);
+    this.settings = new HashMap<>(SERVER_SETTINGS.size());
+    for (Entry<String, PGSetting> entry : SERVER_SETTINGS.entrySet()) {
+      this.settings.put(entry.getKey(), entry.getValue().copy());
+    }
     this.settings.get("server_version").initSettingValue(options.getServerVersion());
+    this.settings.get("server_version_num").initSettingValue(options.getServerVersionNum());
+  }
+
+  /**
+   * Potentially add any session state to the given statement if that is needed. This can for
+   * example include a CTE for pg_settings, if the statement references that table. This method can
+   * be extended in the future to include CTEs for more tables that require session state, or that
+   * are not yet supported in pg_catalog.
+   */
+  public Statement addSessionState(ParsedStatement parsedStatement, Statement statement) {
+    if (parsedStatement.isQuery()
+        && parsedStatement.getSqlWithoutComments().contains("pg_settings")) {
+      String pgSettingsCte = generatePGSettingsCte();
+      // Check whether we can safely use the original SQL statement.
+      String sql =
+          startsWithIgnoreCase(statement.getSql(), "select")
+                  || startsWithIgnoreCase(statement.getSql(), "with")
+              ? statement.getSql()
+              : parsedStatement.getSqlWithoutComments();
+      Statement.Builder builder = null;
+      if (startsWithIgnoreCase(sql, "select")) {
+        builder = Statement.newBuilder("with ").append(pgSettingsCte).append(" ").append(sql);
+      } else if (startsWithIgnoreCase(sql, "with")) {
+        builder =
+            Statement.newBuilder("with ")
+                .append(pgSettingsCte)
+                .append(", ")
+                .append(sql.substring("with".length()));
+      }
+      if (builder != null) {
+        Map<String, Value> parameters = statement.getParameters();
+        for (Entry<String, Value> param : parameters.entrySet()) {
+          builder.bind(param.getKey()).to(param.getValue());
+        }
+        statement = builder.build();
+      }
+    }
+    return statement;
+  }
+
+  static boolean startsWithIgnoreCase(String string, String prefix) {
+    return string.substring(0, prefix.length()).equalsIgnoreCase(prefix);
+  }
+
+  /**
+   * Generates a Common Table Expression that represents the pg_settings table. Note that the
+   * generated query adds two additional CTEs that could in theory hide existing user tables. It is
+   * however strongly recommended that user tables never start with 'pg_', as all system tables in
+   * PostgreSQL start with 'pg_' and 'pg_catalog' is by design always included in the search_path
+   * and is by default the first entry on the search_path. This means that user tables that start
+   * with 'pg_' always risk being hidden by user tables, unless pg_catalog has been explicitly added
+   * to the search_path after one or more user schemas.
+   */
+  String generatePGSettingsCte() {
+    return "pg_settings_inmem_ as (\n"
+        + getAll().stream()
+            .filter(setting -> SUPPORTED_PG_SETTINGS_KEYS.contains(setting.getCasePreservingKey()))
+            .map(PGSetting::getSelectStatement)
+            .collect(Collectors.joining("\nunion all\n"))
+        + "\n),\n"
+        + "pg_settings_names_ as (\n"
+        + "select name from pg_settings_inmem_\n"
+        + "union\n"
+        + "select name from pg_catalog.pg_settings\n"
+        + "),\n"
+        + "pg_settings as (\n"
+        + "select n.name, "
+        + generatePgSettingsColumnExpressions()
+        + "\n"
+        + "from pg_settings_names_ n\n"
+        + "left join pg_settings_inmem_ s1 using (name)\n"
+        + "left join pg_catalog.pg_settings s2 using (name)\n"
+        + "order by name\n"
+        + ")\n";
+  }
+
+  /**
+   * Generates a string of `coalesce(s1.col, s2.col) as col` for all column names (except `name`) in
+   * pg_settings.
+   */
+  private static String generatePgSettingsColumnExpressions() {
+    return PGSetting.getColumnNames().stream()
+        .skip(1)
+        .map(column -> "coalesce(s1." + column + ", s2." + column + ") as " + column)
+        .collect(Collectors.joining(","));
   }
 
   private static String toKey(String extension, String name) {
@@ -59,6 +176,26 @@ private static String toKey(String extension, String name) {
         : extension.toLowerCase(Locale.ROOT) + "." + name.toLowerCase(Locale.ROOT);
   }
 
+  /** Sets the value of the specified setting at connection startup. */
+  public void setConnectionStartupValue(String extension, String name, String value) {
+    String key = toKey(extension, name);
+    PGSetting setting = this.settings.get(key);
+    if (setting == null && extension == null) {
+      // Ignore unknown settings.
+      return;
+    }
+    if (setting == null) {
+      setting = new PGSetting(extension, name);
+      this.settings.put(key, setting);
+    }
+    try {
+      setting.initConnectionValue(value);
+    } catch (Exception ignore) {
+      // ignore errors in startup values to prevent unknown or invalid settings from stopping a
+      // connection from being made.
+    }
+  }
+
   /**
    * Sets the value of the specified setting. The new value will be persisted if the current
    * transaction is committed. The value will be lost if the transaction is rolled back.
@@ -146,7 +283,7 @@ public List<PGSetting> getAll() {
     for (String key : keys) {
       result.add(internalGet(key));
     }
-    result.sort(Comparator.comparing(PGSetting::getKey));
+    result.sort(Comparator.comparing(PGSetting::getCasePreservingKey));
     return result;
   }