Project import generated by Copybara.

README.md

@@ -173,6 +173,9 @@ Notable supported features:
 - The [Cloud Spanner PostgreSQL interface](
   https://cloud.google.com/spanner/docs/postgresql-interface).
 
+- [Connecting with PostgreSQL drivers and tools through PGAdapter](
+  https://github.com/GoogleCloudPlatform/pgadapter/blob/postgresql-dialect/docs/emulator.md).
+
 
 Notable limitations:
 

WORKSPACE

@@ -78,7 +78,7 @@ http_archive(
 
 load("@io_bazel_rules_go//go:deps.bzl", "go_register_toolchains", "go_rules_dependencies")
 
-go_register_toolchains(version = "1.20.12")
+go_register_toolchains(version = "1.20.13")
 
 _bazel_gazelle_version = "0.31.1"
 

backend/common/ids.h

@@ -84,6 +84,13 @@ using SequenceID = std::string;
 // A SequenceID generator. Each database has a single SequenceIDGenerator.
 using SequenceIDGenerator = UniqueIdGenerator<SequenceID>;
 
+// Unique identifier associated with a named schema. NamedSchemaID is guaranteed
+// to be unique within a single database.
+using NamedSchemaID = std::string;
+
+// A NamedSchemaID generator. Each database has a single NamedSchemaIDGenerator.
+using NamedSchemaIDGenerator = UniqueIdGenerator<NamedSchemaID>;
+
 // Unique identifier associated with a transaction.
 using TransactionID = int64_t;
 

backend/query/BUILD

@@ -64,6 +64,7 @@ cc_library(
         "//common:config",
         "//common:constants",
         "//common:errors",
+        "//common:feature_flags",
         "//common:limits",
         "//frontend/converters:values",
         "//third_party/spanner_pg/interface:emulator_parser",

backend/query/analyzer_options.cc

@@ -56,17 +56,14 @@ zetasql::LanguageOptions MakeGoogleSqlLanguageOptions() {
   options.set_name_resolution_mode(zetasql::NAME_RESOLUTION_DEFAULT);
   options.set_product_mode(zetasql::PRODUCT_EXTERNAL);
   options.SetEnabledLanguageFeatures({
-      zetasql::FEATURE_EXTENDED_TYPES,
-      zetasql::FEATURE_NAMED_ARGUMENTS,
-      zetasql::FEATURE_NUMERIC_TYPE,
-      zetasql::FEATURE_TABLESAMPLE,
+      zetasql::FEATURE_EXTENDED_TYPES, zetasql::FEATURE_NAMED_ARGUMENTS,
+      zetasql::FEATURE_NUMERIC_TYPE, zetasql::FEATURE_TABLESAMPLE,
       zetasql::FEATURE_TIMESTAMP_NANOS,
       zetasql::FEATURE_V_1_1_HAVING_IN_AGGREGATE,
       zetasql::FEATURE_V_1_1_NULL_HANDLING_MODIFIER_IN_AGGREGATE,
       zetasql::FEATURE_V_1_1_ORDER_BY_COLLATE,
       zetasql::FEATURE_V_1_1_SELECT_STAR_EXCEPT_REPLACE,
-      zetasql::FEATURE_V_1_2_SAFE_FUNCTION_CALL,
-      zetasql::FEATURE_JSON_TYPE,
+      zetasql::FEATURE_V_1_2_SAFE_FUNCTION_CALL, zetasql::FEATURE_JSON_TYPE,
       zetasql::FEATURE_JSON_ARRAY_FUNCTIONS,
       zetasql::FEATURE_JSON_STRICT_NUMBER_PARSING,
       zetasql::FEATURE_V_1_3_DML_RETURNING,

backend/query/catalog_test.cc

@@ -230,6 +230,9 @@ TEST_F(CatalogTest, FindBuiltInTableValuedFunction) {
   EXPECT_EQ(tvf->FullName(), "ML.PREDICT");
   ZETASQL_EXPECT_OK(catalog().FindTableValuedFunction({"SAFE.ML.PREDICT"}, &tvf, {}));
   EXPECT_EQ(tvf->FullName(), "SAFE.ML.PREDICT");
+  ZETASQL_EXPECT_OK(
+      catalog().FindTableValuedFunction({"pg.jsonb_array_elements"}, &tvf, {}));
+  EXPECT_EQ(tvf->FullName(), "pg.jsonb_array_elements");
 }
 
 TEST_F(CatalogTest, FindTableValuedFunction) {

backend/query/function_catalog.cc

@@ -56,7 +56,9 @@ namespace backend {
 
 namespace {
 using postgres_translator::GetSpannerPGFunctions;
+using postgres_translator::GetSpannerPGTVFs;
 using postgres_translator::SpannerPGFunctions;
+using postgres_translator::SpannerPGTVFs;
 
 absl::StatusOr<zetasql::Value> EvalPendingCommitTimestamp(
     absl::Span<const zetasql::Value> args) {
@@ -197,6 +199,12 @@ void FunctionCatalog::AddSpannerPGFunctions(
   for (auto& function : spanner_pg_functions) {
     functions_[function->Name()] = std::move(function);
   }
+
+  SpannerPGTVFs spanner_pg_tvfs = GetSpannerPGTVFs(catalog_name_);
+
+  for (auto& tvf : spanner_pg_tvfs) {
+    table_valued_functions_[tvf->FullName()] = std::move(tvf);
+  }
 }
 
 void FunctionCatalog::GetFunction(const std::string& name,
@@ -261,9 +269,16 @@ FunctionCatalog::GetInternalSequenceStateFunction(
     if (latest_schema_ == nullptr) {
       return error::SequenceNeedsAccessToSchema();
     }
-    // ZetaSQL algebrizer prepends a prefix to the sequence name.
-    std::string sequence_name =
-        std::string(absl::StripPrefix(args[0].string_value(), "_sequence_"));
+
+    std::string sequence_name;
+    if (latest_schema_->dialect() ==
+        database_api::DatabaseDialect::POSTGRESQL) {
+      sequence_name = args[0].string_value();
+    } else {
+      // ZetaSQL algebrizer prepends a prefix to the sequence name.
+      sequence_name =
+          std::string(absl::StripPrefix(args[0].string_value(), "_sequence_"));
+    }
     const backend::Sequence* sequence =
         latest_schema_->FindSequence(sequence_name);
     if (sequence == nullptr) {
@@ -305,9 +320,16 @@ FunctionCatalog::GetNextSequenceValueFunction(const std::string& catalog_name) {
     if (latest_schema_ == nullptr) {
       return error::SequenceNeedsAccessToSchema();
     }
+
     // ZetaSQL algebrizer prepends a prefix to the sequence name.
-    std::string sequence_name =
-        std::string(absl::StripPrefix(args[0].string_value(), "_sequence_"));
+    std::string sequence_name;
+    if (latest_schema_->dialect() ==
+        database_api::DatabaseDialect::POSTGRESQL) {
+      sequence_name = args[0].string_value();
+    } else {
+      sequence_name =
+          std::string(absl::StripPrefix(args[0].string_value(), "_sequence_"));
+    }
     const backend::Sequence* sequence =
         latest_schema_->FindSequence(sequence_name);
     if (sequence == nullptr) {

backend/query/function_catalog.h

@@ -33,8 +33,7 @@ namespace spanner {
 namespace emulator {
 namespace backend {
 
-constexpr char kCloudSpannerEmulatorFunctionCatalogName[] =
-    "CloudSpannerEmulatorCatalog";
+constexpr char kCloudSpannerEmulatorFunctionCatalogName[] = "Spanner";
 
 // A catalog of all SQL functions.
 //

backend/query/index_hint_validator.h

@@ -33,9 +33,11 @@ namespace backend {
 class IndexHintValidator : public zetasql::ResolvedASTVisitor {
  public:
   IndexHintValidator(const Schema* schema,
-                     bool disable_null_filtered_index_check = false)
+                     bool disable_null_filtered_index_check = false
+                     )
       : schema_(schema),
-        disable_null_filtered_index_check_(disable_null_filtered_index_check) {}
+        disable_null_filtered_index_check_(disable_null_filtered_index_check)
+  {}
 
  private:
   absl::Status VisitResolvedQueryStmt(

backend/query/partitionability_validator.cc

@@ -33,9 +33,47 @@ namespace spanner {
 namespace emulator {
 namespace backend {
 
+namespace {
+
+class HasTableScanVisitor : public zetasql::ResolvedASTVisitor {
+ public:
+  HasTableScanVisitor() = default;
+
+  absl::StatusOr<bool> HasTableScan(const zetasql::ResolvedNode* node) {
+    has_table_scan_ = false;
+    ZETASQL_RETURN_IF_ERROR(node->Accept(this));
+    return has_table_scan_;
+  }
+
+  absl::Status VisitResolvedTableScan(
+      const zetasql::ResolvedTableScan* node) override {
+    has_table_scan_ = true;
+    return absl::OkStatus();
+  }
+
+ private:
+  bool has_table_scan_ = false;
+};
+
+// Returns true if the expression tree has TableScan.
+absl::StatusOr<bool> HasTableScan(const zetasql::ResolvedNode* node) {
+  HasTableScanVisitor visitor;
+  return visitor.HasTableScan(node);
+}
+
+}  // namespace
+
 absl::Status PartitionabilityValidator::ValidatePartitionability(
     const zetasql::ResolvedNode* node) {
   ABSL_DCHECK(node->node_kind() == zetasql::RESOLVED_QUERY_STMT);
+  if (EmulatorFeatureFlags::instance()
+          .flags()
+          .enable_batch_query_with_no_table_scan) {
+    ZETASQL_ASSIGN_OR_RETURN(bool has_table_scan, HasTableScan(node));
+    if (!has_table_scan) {
+      return absl::OkStatus();
+    }
+  }
   if (HasSubquery(node)) {
     return error::NonPartitionableQuery("Query contains subquery.");
   }

backend/query/partitionability_validator_test.cc

@@ -160,28 +160,41 @@ TEST_F(PartitionabilityValidatorTest,
               zetasql_base::testing::StatusIs(absl::StatusCode::kInvalidArgument));
 }
 
-TEST_F(PartitionabilityValidatorTest, ValidateEmptyQueryNotPartitionable) {
+TEST_F(PartitionabilityValidatorTest, SelectOneIsPartitionable) {
+  test::ScopedEmulatorFeatureFlagsSetter feature_flag_setter(
+      EmulatorFeatureFlags::Flags{.enable_batch_query_with_no_table_scan =
+                                      true});
   PartitionabilityValidator validator{schema()};
-  std::unique_ptr<const zetasql::ResolvedQueryStmt> query_stmt =
-      zetasql::MakeResolvedQueryStmt();
 
-  ASSERT_THAT(query_stmt->Accept(&validator),
-              zetasql_base::testing::StatusIs(absl::StatusCode::kInvalidArgument));
+  std::unique_ptr<const zetasql::AnalyzerOutput> analyzer_output;
+  zetasql::SimpleCatalog catalog("test simple catalog");
+  zetasql::TypeFactory type_factory;
+  ZETASQL_ASSERT_OK(zetasql::AnalyzeStatement("SELECT 1",
+                                        zetasql::AnalyzerOptions(), &catalog,
+                                        &type_factory, &analyzer_output));
+
+  const zetasql::ResolvedStatement* resolved_statement =
+      analyzer_output->resolved_statement();
+  ZETASQL_EXPECT_OK(resolved_statement->Accept(&validator));
 }
 
-TEST_F(PartitionabilityValidatorTest, ValidateNoTableScanNonPartitionable) {
+TEST_F(PartitionabilityValidatorTest,
+       SelectConstValuesWithSortIsPartitionable) {
+  test::ScopedEmulatorFeatureFlagsSetter feature_flag_setter(
+      EmulatorFeatureFlags::Flags{.enable_batch_query_with_no_table_scan =
+                                      true});
   PartitionabilityValidator validator{schema()};
-  QueryableTable table{schema()->FindTable("test_table"), /*reader=*/nullptr};
 
-  std::unique_ptr<const zetasql::ResolvedScan> project_scan =
-      zetasql::MakeResolvedProjectScan();
-  std::unique_ptr<const zetasql::ResolvedQueryStmt> query_stmt =
-      zetasql::MakeResolvedQueryStmt(/*output_column_list=*/{},
-                                       /*is_value_table=*/false,
-                                       std::move(project_scan));
+  std::unique_ptr<const zetasql::AnalyzerOutput> analyzer_output;
+  zetasql::SimpleCatalog catalog("test simple catalog");
+  zetasql::TypeFactory type_factory;
+  ZETASQL_ASSERT_OK(zetasql::AnalyzeStatement(
+      "SELECT a FROM UNNEST([3, 2, 1]) AS a ORDER BY a",
+      zetasql::AnalyzerOptions(), &catalog, &type_factory, &analyzer_output));
 
-  ASSERT_THAT(query_stmt->Accept(&validator),
-              zetasql_base::testing::StatusIs(absl::StatusCode::kInvalidArgument));
+  const zetasql::ResolvedStatement* resolved_statement =
+      analyzer_output->resolved_statement();
+  ZETASQL_EXPECT_OK(resolved_statement->Accept(&validator));
 }
 
 }  // namespace

backend/query/query_engine.cc

@@ -69,6 +69,7 @@
 #include "common/config.h"
 #include "common/constants.h"
 #include "common/errors.h"
+#include "common/feature_flags.h"
 #include "common/limits.h"
 #include "frontend/converters/values.h"
 #include "third_party/spanner_pg/interface/emulator_parser.h"
@@ -410,6 +411,19 @@ absl::StatusOr<ExecuteUpdateResult> EvaluateResolvedInsert(
     const zetasql::ResolvedInsertStmt* insert_statement,
     const zetasql::ParameterValueMap& parameters,
     zetasql::TypeFactory* type_factory) {
+  bool is_upsert_query = insert_statement->insert_mode() ==
+                             zetasql::ResolvedInsertStmt::OR_IGNORE ||
+                         insert_statement->insert_mode() ==
+                             zetasql::ResolvedInsertStmt::OR_UPDATE;
+  if (is_upsert_query &&
+      !EmulatorFeatureFlags::instance().flags().enable_upsert_queries) {
+    std::string insert_mode_name = (insert_statement->insert_mode() ==
+                                    zetasql::ResolvedInsertStmt::OR_IGNORE)
+                                       ? "Insert or ignore"
+                                       : "Insert or update";
+    return error::UnsupportedUpsertQueries(insert_mode_name);
+  }
+
   ZETASQL_ASSIGN_OR_RETURN(auto pending_ts_columns,
                    PendingCommitTimestampColumnsInInsert(
                        insert_statement->insert_column_list(),
@@ -663,8 +677,10 @@ ExtractValidatedResolvedStatementAndOptions(
 
   // Validate the index hints.
   IndexHintValidator index_hint_validator{
-      context.schema, options.disable_query_null_filtered_index_check ||
-                          config::disable_query_null_filtered_index_check()};
+      context.schema,
+      options.disable_query_null_filtered_index_check ||
+          config::disable_query_null_filtered_index_check()
+  };
   ZETASQL_RETURN_IF_ERROR(statement->Accept(&index_hint_validator));
 
   // Check the query size limits

backend/query/query_engine_test.cc

@@ -737,21 +737,15 @@ TEST_P(QueryEngineTest, PartitionableSimpleScanSubqueryColumn) {
 }
 
 TEST_P(QueryEngineTest, PartitionableSimpleScanNoTable) {
-  std::string error_msg = kQueryNotASimpleTableScanError;
   Query query{"SELECT a FROM UNNEST(ARRAY[1, 2, 3]) AS a"};
-  EXPECT_THAT(
-      query_engine().IsPartitionable(
-          query, QueryContext{multi_table_schema(), reader()}),
-      StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr(error_msg)));
+  ZETASQL_ASSERT_OK(query_engine().IsPartitionable(
+      query, QueryContext{multi_table_schema(), reader()}));
 }
 
 TEST_P(QueryEngineTest, PartitionableSimpleScanFilterNoTable) {
-  std::string error_msg = kQueryNotASimpleTableScanError;
   Query query{"SELECT a FROM UNNEST(ARRAY[1, 2, 3]) AS a WHERE a = 1"};
-  EXPECT_THAT(
-      query_engine().IsPartitionable(
-          query, QueryContext{multi_table_schema(), reader()}),
-      StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr(error_msg)));
+  ZETASQL_ASSERT_OK(query_engine().IsPartitionable(
+      query, QueryContext{multi_table_schema(), reader()}));
 }
 
 TEST_P(QueryEngineTest, PartitionableExecuteSqlSimpleScanFilterSubquery) {

backend/query/query_validator.cc

@@ -30,6 +30,8 @@
 #include "absl/strings/string_view.h"
 #include "backend/query/feature_filter/gsql_supported_functions.h"
 #include "backend/query/feature_filter/sql_feature_filter.h"
+#include "backend/query/query_context.h"
+#include "backend/query/query_engine_options.h"
 #include "backend/schema/catalog/index.h"
 #include "backend/schema/catalog/sequence.h"
 #include "common/constants.h"
@@ -273,7 +275,9 @@ absl::Status QueryValidator::CheckHintValue(
       }
       const Index* index = context_.schema->FindIndex(index_name);
       if (index == nullptr) {
-        return error::InvalidHintValue(name, value.DebugString());
+        // We don't have the table name here. So this will not match prod error
+        // message.
+        return error::QueryHintIndexNotFound("", index_name);
       }
 
       indexes_used_.insert(index);

backend/schema/backfills/BUILD

@@ -92,15 +92,21 @@ cc_test(
     ],
     deps = [
         ":schema_backfillers",
+        "//backend/access:read",
+        "//backend/access:write",
         "//backend/database",
+        "//backend/datamodel:value",
         "//backend/schema/catalog:schema",
         "//backend/schema/updater:schema_updater",
         "//backend/transaction:read_only_transaction",
+        "//backend/transaction:read_write_transaction",
+        "//common:clock",
         "//common:errors",
         "//tests/common:proto_matchers",
         "@com_github_grpc_grpc//:grpc++",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",
         "@com_google_zetasql//zetasql/base/testing:status_matchers",