main.cpp

/**
 * Copyright 2023 KUMAZAKI Hiroki
 *
 * 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.
 */

#include <algorithm>
#include <exception>
#include <iostream>
#include <iterator>
#include <sstream>
#include <string>
#include <utility>
#include <vector>

#include "common/constants.hpp"
#include "common/status_or.hpp"
#include "database/database.hpp"
#include "executor/executor_base.hpp"
#include "query/sql_engine.hpp"
#include "recovery/recovery_manager.hpp"
#include "server/postgres_protocol.hpp"
#include "type/row.hpp"

int main(int argc, char** argv) {
  // --force: a torn/unparsable WAL tail is truncated to its intact prefix
  // instead of treating the corruption as fatal (the default).
  bool force = false;
  std::vector<std::string> args(argv + 1, argv + argc);
  auto flag = std::remove(args.begin(), args.end(), "--force");
  if (flag != args.end()) {
    force = true;
    args.erase(flag, args.end());
  }
  if (args.size() != 1) {
    std::cerr << "usage: tinylamb [--force] <database-file> < input.sql\n";
    return 2;
  }
  tinylamb::RecoveryManager::SetTornTailTruncationAllowed(force);

  std::string sql((std::istreambuf_iterator<char>(std::cin)),
                  std::istreambuf_iterator<char>());
  const std::vector<std::string> statements =
      tinylamb::pgwire::SplitSqlStatements(sql);
  if (statements.empty()) {
    std::cerr << "no SQL was provided on standard input\n";
    return 2;
  }

  auto database_or_status = tinylamb::Database::Create(args[0]);
  if (!database_or_status.HasValue()) {
    std::cerr << "failed to open database " << args[0] << ": "
              << database_or_status.GetStatus() << '\n';
    return 1;
  }
  tinylamb::Database& database = *database_or_status.Value();
  tinylamb::TransactionContext context = database.BeginContext();
  tinylamb::SqlEngine engine(database);
  // One implicit transaction wraps every statement of the script; the first
  // failure aborts the whole run instead of terminating the process.
  // D4 (docs/design.md): output is buffered and only printed after the
  // commit's durability barrier, so no row reaches the user before the
  // commits it observed are durable.
  std::vector<std::string> pending_output;
  for (const std::string& statement : statements) {
    try {
      tinylamb::StatusOr<tinylamb::QueryResult> executed =
          engine.Execute(context, statement);
      if (!executed.HasValue()) {
        const std::string& last_error = engine.LastError();
        std::cerr << "SQL error: " << last_error;
        if (last_error.empty()) {
          std::cerr << executed.GetStatus();
        }
        std::cerr << '\n';
        context.Abort();
        return 1;
      }
      std::string buffer;
      executed.Value().ForEach([&buffer](const tinylamb::Row& row) {
        std::ostringstream line;
        line << row << '\n';
        buffer += line.str();
      });
      // Executors latch drain-time failures on their status channel
      // (Next()==false reads as EOF): an unchecked overflow would print
      // as a successful empty result and exit 0.
      if (const tinylamb::Status drain_status = executed.Value().GetStatus();
          drain_status != tinylamb::Status::kSuccess) {
        std::cerr << "SQL error: " << drain_status << '\n';
        context.Abort();
        return 1;
      }
      pending_output.push_back(std::move(buffer));
    } catch (const std::exception& error) {
      std::cerr << "error executing statement: " << error.what() << '\n';
      context.Abort();
      return 1;
    }
  }
  try {
    if (context.PreCommit() != tinylamb::Status::kSuccess) {
      // Keep the cleanup symmetric with the Prepare failure path above.
      std::cerr << "transaction commit failed\n";
      context.Abort();
      return 1;
    }
  } catch (const std::exception& error) {
    // TransactionManager::PreCommit already released the locks and marked the
    // transaction aborted before rethrowing, so no extra Abort() here.
    std::cerr << "error during commit: " << error.what() << '\n';
    return 1;
  }
  // Durability barrier passed; the buffered rows may now reach the user.
  for (const std::string& chunk : pending_output) {
    std::cout << chunk;
  }
  std::cout.flush();
  return 0;
}