// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef FEDERATED_EXAMPLE_DATABASE_H_
#define FEDERATED_EXAMPLE_DATABASE_H_

#include <cstdint>

#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>

#include <absl/status/statusor.h>
#include <base/files/file_path.h>
#include <base/time/time.h>
#include <sqlite3.h>

namespace federated {

// Example objects stored in corresponding `client_name` tables.
// An example represents a training example of federated computation.
struct ExampleRecord {
  // The ID of this example in the client table. Only populated in records
  // being retrieved from (c.f. being inserted into) the example database.
  int64_t id = -1;

  std::string serialized_example;
  base::Time timestamp;
};

// Provides access to example database.
//
// WARNING: Do not pass strings to these methods (e.g. Init, GetIterator) that
//          have not been carefully sanitized. This class does not perform
//          string sanitization and is therefore susceptible to SQL code
//          injection.
//
// Example usage:
// Construct and initialize:
//    ExampleDatabase db(db_path);
//    if(!db.Init(kTestClients) || !db.IsOpen() || !db.CheckIntegrity()) {
//      // Error handling.
//    }
//
// Insert an example:
//    ExampleRecord example_record;
//    example_record.serialized_example = serialized_example;
//    example_record.timestamp = base::Time::Now();
//    db.InsertExample(client_name, example_record);
//
// Query examples:
//    ExampleDatabase::Iterator it = db.GetIterator("client_1");
//    while (true) {
//      const absl::StatusOr<ExampleRecord> result = it.Next();
//      if (result.ok()) {
//        // Do something with *result.
//
//        continue;
//      }
//
//      if (absl::IsOutOfRange(result)) {
//        // End of iterator.
//      } else {
//        // Handle error.
//      }
//    }
//
// See example_database_test.cc and storage_manager_impl.cc for more details.

class ExampleDatabase {
 public:
  // Handles one read-only iteration through a table.
  struct Iterator final {
   public:
    Iterator(sqlite3* db, const std::string& client_name);
    Iterator(Iterator&& other);
    Iterator(const Iterator& other) = delete;
    Iterator& operator=(const Iterator& other) = delete;
    ~Iterator();

    // Returns the next example, an "out of range" error if the end of the
    // iteration has been reached, or any other error if example fetching
    // failed.
    absl::StatusOr<ExampleRecord> Next();

    // Releases sqlite resources / locks.
    //
    // Called automatically when the iteration is complete or the iterator is
    // destroyed, but must be called manually when iteration is abandoned
    // early. The database cannot be closed unless all iterators have been
    // closed by one means or another.
    void Close();

   private:
    sqlite3_stmt* stmt_;
  };

  // Creates an instance to talk to the database file at `db_path`. Init() must
  // be called to establish connection.
  explicit ExampleDatabase(const base::FilePath& db_path);
  ExampleDatabase(const ExampleDatabase&) = delete;
  ExampleDatabase& operator=(const ExampleDatabase&) = delete;

  virtual ~ExampleDatabase();

  // Initializes database connection. Must be called before any other queries.
  // Returns true if no error occurred.
  //
  // WARNING: client names are used to construct SQL statements but are not
  //          sanitized in any way. Therefore this method is susceptible to
  //          code injection unless the provided names are carefully vetted or
  //          sanitized.
  virtual bool Init(const std::unordered_set<std::string>& clients);
  // Returns true if the database connection is open.
  virtual bool IsOpen() const;
  // Closes database connection. Returns true if no error occurred.
  virtual bool Close();
  // Runs sqlite built-in integrity check. Returns true if no error is found.
  virtual bool CheckIntegrity() const;

  // Returns an iterator through the examples for the given client.
  //
  // WARNING: client names are used to construct SQL statements but are not
  //          sanitized in any way. Therefore this method is susceptible to
  //          code injection unless the provided names are carefully vetted or
  //          sanitized.
  virtual Iterator GetIterator(const std::string& client_name) const;

  // Inserts example into the table matching its client_name. Returns true
  // if no error occurred.
  //
  // WARNING: client names are used to construct SQL statements but are not
  //          sanitized in any way. Therefore this method is susceptible to
  //          code injection unless the provided names are carefully vetted or
  //          sanitized.
  virtual bool InsertExample(const std::string& client_name,
                             const ExampleRecord& example_record);

  // Returns the count of examples in the client's table.
  //
  // WARNING: client names are used to construct SQL statements but are not
  //          sanitized in any way. Therefore this method is susceptible to
  //          code injection unless the provided names are carefully vetted or
  //          sanitized.
  virtual int ExampleCount(const std::string& client_name) const;

  // Deletes all examples in the specified client table. We expose only this
  // rudimentary functionality since small federated clients typically delete
  // all training examples at the end of a training session. More sophiscated
  // behavior can be added if it is later needed.
  //
  // WARNING: client names are used to construct SQL statements but are not
  //          sanitized in any way. Therefore this method is susceptible to
  //          code injection unless the provided names are carefully vetted or
  //          sanitized.
  virtual void DeleteAllExamples(const std::string& client_name);

  sqlite3* sqlite3_for_testing() const { return db_.get(); }

 private:
  // Typedef of sqlite3_exec callback, see sqlite doc:
  // https://sqlite.org/c3ref/exec.html.
  using SqliteCallback = int (*)(void* /*data*/,
                                 int /*count*/,
                                 char** /*row*/,
                                 char** /*names*/);

  // Sqlite error code and error message.
  struct ExecResult {
    int code;
    std::string error_msg;
  };

  // Returns true if the client's table exists.
  bool ClientTableExists(const std::string& client_name) const;

  // Returns true if the client's table is created without error.
  bool CreateClientTable(const std::string& client_name);

  // Executes sql.
  ExecResult ExecSql(const std::string& sql) const;
  ExecResult ExecSql(const std::string& sql,
                     SqliteCallback callback,
                     void* data) const;

  const base::FilePath db_path_;
  std::unique_ptr<sqlite3, decltype(&sqlite3_close)> db_;
};

}  // namespace federated

#endif  // FEDERATED_EXAMPLE_DATABASE_H_