engine/schema/src/main/java/com/cloud/upgrade/dao/Upgrade227to228Premium.java

// 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 com.cloud.upgrade.dao;

import java.io.InputStream;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;

import org.apache.log4j.Logger;

import com.cloud.utils.exception.CloudRuntimeException;

public class Upgrade227to228Premium extends Upgrade227to228 {
    final static Logger s_logger = Logger.getLogger(Upgrade227to228Premium.class);

    @Override
    public InputStream[] getPrepareScripts() {
        InputStream[] newScripts = new InputStream[2];
        newScripts[0] = super.getPrepareScripts()[0];
        final String scriptFile = "META-INF/db/schema-227to228-premium.sql";
        final InputStream script = Thread.currentThread().getContextClassLoader().getResourceAsStream(scriptFile);
        if (script == null) {
            throw new CloudRuntimeException("Unable to find " + scriptFile);
        }
        newScripts[1] = script;

        return newScripts;
    }

    @Override
    public void performDataMigration(Connection conn) {
        addSourceIdColumn(conn);
        addNetworkIdsToUserStats(conn);
        super.performDataMigration(conn);
    }

    private void addSourceIdColumn(Connection conn) {
        boolean insertField = false;
        try {
            PreparedStatement pstmt;
            try {
                pstmt = conn.prepareStatement("SELECT source_id FROM `cloud_usage`.`usage_storage`");
                ResultSet rs = pstmt.executeQuery();

                if (rs.next()) {
                    s_logger.info("The source id field already exist, not adding it");
                }

            } catch (Exception e) {
                // if there is an exception, it means that field doesn't exist, and we can create it
                insertField = true;
            }

            if (insertField) {
                s_logger.debug("Adding source_id to usage_storage...");
                pstmt = conn.prepareStatement("ALTER TABLE `cloud_usage`.`usage_storage` ADD COLUMN `source_id` bigint unsigned");
                pstmt.executeUpdate();
                s_logger.debug("Column source_id was added successfully to usage_storage table");
                pstmt.close();
            }

        } catch (SQLException e) {
            s_logger.error("Failed to add source_id to usage_storage due to ", e);
            throw new CloudRuntimeException("Failed to add source_id to usage_storage due to ", e);
        }
    }

    private void addNetworkIdsToUserStats(Connection conn) {
        s_logger.debug("Adding network IDs to user stats...");
        try {
            String stmt = "SELECT DISTINCT public_ip_address FROM `cloud`.`user_statistics` WHERE public_ip_address IS NOT null";
            PreparedStatement pstmt = conn.prepareStatement(stmt);
            ResultSet rs = pstmt.executeQuery();

            while (rs.next()) {
                String publicIpAddress = rs.getString(1);
                stmt = "SELECT network_id FROM `cloud`.`user_ip_address` WHERE public_ip_address = ?";
                pstmt = conn.prepareStatement(stmt);
                pstmt.setString(1, publicIpAddress);
                ResultSet rs2 = pstmt.executeQuery();

                if (rs2.next()) {
                    Long networkId = rs2.getLong(1);
                    String[] dbs = {"cloud", "cloud_usage"};
                    for (String db : dbs) {
                        stmt = "UPDATE `" + db + "`.`user_statistics` SET network_id = ? WHERE public_ip_address = ?";
                        pstmt = conn.prepareStatement(stmt);
                        pstmt.setLong(1, networkId);
                        pstmt.setString(2, publicIpAddress);
                        pstmt.executeUpdate();
                    }
                }

                rs2.close();
            }

            rs.close();
            pstmt.close();
            s_logger.debug("Successfully added network IDs to user stats.");
        } catch (SQLException e) {
            String errorMsg = "Failed to add network IDs to user stats.";
            s_logger.error(errorMsg, e);
            throw new CloudRuntimeException(errorMsg, e);
        }
    }

}