content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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 org.apache.beam.sdk.extensions.sql.meta.store;
  
  import java.util.HashMap;
  import java.util.Map;
  import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTable;
  import org.apache.beam.sdk.extensions.sql.meta.Table;
  import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  /**
   * A {@link MetaStore} which stores the meta info in memory.
   *
   * <p>NOTE, because this implementation is memory based, the metadata is NOT persistent. for tables
   * which created, you need to create again every time you launch the {@link
   * org.apache.beam.sdk.extensions.sql.BeamSqlCli}.
   */
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class InMemoryMetaStore implements MetaStore {
    private Map<String, Table> tables = new HashMap<>();
    private Map<String, TableProvider> providers = new HashMap<>();
  
    @Override
    public String getTableType() {
      return "store";
    }
  
    @Override
    public void createTable(Table table) {
      validateTableType(table);
  
      // first assert the table name is unique
      if (tables.containsKey(table.getName())) {
        throw new IllegalArgumentException("Duplicate table name: " + table.getName());
      }
  
      // invoke the provider's create
      getProvider(table.getType()).createTable(table);
  
      // store to the global metastore
      tables.put(table.getName(), table);
    }
  
    @Override
    public void dropTable(String tableName) {
      if (!tables.containsKey(tableName)) {
        throw new IllegalArgumentException("No such table: " + tableName);
      }
  
      Table table = tables.get(tableName);
      getProvider(table.getType()).dropTable(tableName);
      tables.remove(tableName);
    }
  
    @Override
    public Map<String, Table> getTables() {
      return ImmutableMap.copyOf(tables);
    }
  
    @Override
    public BeamSqlTable buildBeamSqlTable(Table table) {
      TableProvider provider = getProvider(table.getType());
  
      return provider.buildBeamSqlTable(table);
    }
  
    protected void validateTableType(Table table) {
      if (providers.containsKey(table.getType().toLowerCase())) {
        return;
      }
      // check if there is a nested metastore that supports this table
      @Nullable
      InMemoryMetaStore nestedMemoryMetastore = (InMemoryMetaStore) providers.get(getTableType());
      if (nestedMemoryMetastore != null) {
        nestedMemoryMetastore.validateTableType(table);
      } else {
        throw new IllegalArgumentException("Table type: " + table.getType() + " not supported!");
      }
    }
  
    @Override
    public void registerProvider(TableProvider provider) {
      String type = provider.getTableType().toLowerCase();
      if (providers.containsKey(type)) {
        throw new IllegalArgumentException("Provider is already registered for table type: " + type);
      }
  
      initTablesFromProvider(provider);
      this.providers.put(type, provider);
    }
  
    private void initTablesFromProvider(TableProvider provider) {
      Map<String, Table> tables = provider.getTables();
      for (String tableName : tables.keySet()) {
        if (this.tables.containsKey(tableName)) {
          throw new IllegalStateException(
              "Duplicate table: " + tableName + " from provider: " + provider);
        }
      }
      this.tables.putAll(tables);
    }
  
    @Override
    public Map<String, TableProvider> tableProviders() {
      return providers;
    }
  
    @Override
    public boolean supportsPartitioning(Table table) {
      return getProvider(table.getType()).supportsPartitioning(table);
    }
  
    /**
     * Fetches a {@link TableProvider} for this type. This provider can exist in the current {@link
     * InMemoryMetaStore} or a nested {@link InMemoryMetaStore}.
     *
     * @param type
     * @return
     */
    public TableProvider getProvider(String type) {
      @Nullable TableProvider provider = providers.get(type.toLowerCase());
      if (provider != null) {
        return provider;
      }
  
      // check nested InMemoryMetaStore
      provider = providers.get(getTableType());
      if (provider != null && (provider instanceof InMemoryMetaStore)) {
        return ((InMemoryMetaStore) provider).getProvider(type);
      }
  
      throw new IllegalStateException("No TableProvider registered for table type: " + type);
    }
  }
  