public void execute()

in baremaps-core/src/main/java/org/apache/baremaps/tasks/ImportGeoParquet.java [72:176]
75 lines of code
12 McCabe index (conditional complexity)

  public void execute(WorkflowContext context) throws Exception {
    // Validate required parameters
    if (uri == null) {
      throw new WorkflowException("GeoParquet URI cannot be null");
    }
    if (tableName == null || tableName.isEmpty()) {
      throw new WorkflowException("Table name cannot be null or empty");
    }
    if (database == null) {
      throw new WorkflowException("Database connection cannot be null");
    }
    if (databaseSrid == null) {
      throw new WorkflowException("Target SRID cannot be null");
    }

    logger.info("Importing GeoParquet from: {}", uri);

    var dataSource = context.getDataSource(database);

    // Sanitize table name to prevent SQL injection
    String sanitizedTableName = sanitizeTableName(tableName);
    logger.info("Creating table: {}", sanitizedTableName);

    // Set ThreadLocal DataSource for PostgresDdlExecutor to use
    PostgresDdlExecutor.setThreadLocalDataSource(dataSource);

    try {
      // Setup Calcite connection properties
      Properties info = new Properties();
      info.setProperty("lex", "MYSQL");
      info.setProperty("caseSensitive", "false");
      info.setProperty("unquotedCasing", "TO_LOWER");
      info.setProperty("quotedCasing", "TO_LOWER");
      info.setProperty("parserFactory", PostgresDdlExecutor.class.getName() + "#PARSER_FACTORY");

      // Create a connection to Calcite
      try (Connection connection = DriverManager.getConnection("jdbc:calcite:", info)) {

        // Get the list of tables in the GeoParquet
        String[] tables = getGeoParquetTables(connection);

        if (tables.length == 0) {
          logger.warn("No tables found in GeoParquet: {}", uri);
          return;
        }

        // Import each table
        for (String sourceTableName : tables) {
          // Create a temporary table name for the GeoParquet data
          String tempTableName =
              "geoparquet_data_" + System.currentTimeMillis() + "_" + sourceTableName;

          // Register the GeoParquet table in the Calcite schema
          String registerSql = "CREATE TABLE " + tempTableName + " AS " +
              "SELECT * FROM " + sourceTableName;

          logger.info("Executing SQL: {}", registerSql);

          // Execute the DDL statement to create the table
          try (Statement statement = connection.createStatement()) {
            statement.execute(registerSql);
          }

          // Set SRID on geometry column if specified
          try (Connection pgConnection = dataSource.getConnection();
              Statement stmt = pgConnection.createStatement()) {
            stmt.execute(String.format(
                "SELECT UpdateGeometrySRID('%s', 'geometry', %d)",
                sanitizedTableName, databaseSrid));
          }

          // Verify that the table was created in PostgreSQL
          try (Connection pgConnection = dataSource.getConnection();
              Statement statement = pgConnection.createStatement();
              ResultSet resultSet = statement.executeQuery(
                  "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = '" +
                      sanitizedTableName + "')")) {
            if (!resultSet.next() || !resultSet.getBoolean(1)) {
              throw new WorkflowException("Failed to create table: " + sanitizedTableName);
            }
          }

          // Verify that the table has data
          try (Connection pgConnection = dataSource.getConnection();
              Statement statement = pgConnection.createStatement();
              ResultSet resultSet = statement.executeQuery(
                  "SELECT COUNT(*) FROM " + sanitizedTableName)) {
            if (resultSet.next()) {
              int count = resultSet.getInt(1);
              logger.info("Imported {} rows to table: {}", count, sanitizedTableName);
              if (count == 0) {
                logger.warn("No rows were imported from GeoParquet to table: {}",
                    sanitizedTableName);
              }
            }
          }
        }
      }
    } finally {
      // Clean up thread local storage
      PostgresDdlExecutor.clearThreadLocalDataSource();
    }

    logger.info("Successfully imported GeoParquet to table: {}", sanitizedTableName);
  }