<commit_before>package com.linkedin.uif.source.extractor.extract.restapi;

import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.http.HttpEntity;
import org.apache.http.HttpHost;
import org.apache.http.HttpResponse;
import org.apache.http.StatusLine;
import org.apache.http.client.HttpClient;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpRequestBase;
import org.apache.http.conn.params.ConnRoutePNames;
import org.apache.http.impl.client.DefaultHttpClient;
import org.apache.http.util.EntityUtils;

import com.google.common.base.Joiner;
import com.google.common.base.Strings;
import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;

import com.linkedin.uif.configuration.ConfigurationKeys;
import com.linkedin.uif.configuration.WorkUnitState;
import com.linkedin.uif.source.extractor.watermark.Predicate;
import com.linkedin.uif.source.extractor.DataRecordException;
import com.linkedin.uif.source.extractor.exception.HighWatermarkException;
import com.linkedin.uif.source.extractor.exception.RecordCountException;
import com.linkedin.uif.source.extractor.exception.RestApiConnectionException;
import com.linkedin.uif.source.extractor.exception.RestApiProcessingException;
import com.linkedin.uif.source.extractor.exception.SchemaException;
import com.linkedin.uif.source.extractor.extract.BaseExtractor;
import com.linkedin.uif.source.extractor.extract.Command;
import com.linkedin.uif.source.extractor.extract.CommandOutput;
import com.linkedin.uif.source.extractor.extract.SourceSpecificLayer;
import com.linkedin.uif.source.extractor.schema.Schema;
import com.linkedin.uif.source.extractor.utils.Utils;
import com.linkedin.uif.source.workunit.WorkUnit;

/**
 * An implementation of rest api extractor for the sources that are using rest api
 *
 * @param <D> type of data record
 * @param <S> type of schema
 */
public abstract class RestApiExtractor extends BaseExtractor<JsonArray, JsonElement> implements SourceSpecificLayer<JsonArray, JsonElement>, RestApiSpecificLayer {
	private static final Gson gson = new Gson();
	private HttpClient httpClient = null;
	private boolean autoEstablishAuthToken = false;
	private long authTokenTimeout;
	private String accessToken;
	private long createdAt;
	protected String instanceUrl;
	protected String updatedQuery;
	protected Logger log = LoggerFactory.getLogger(RestApiExtractor.class);

	RestApiExtractor(WorkUnitState state) {
		super(state);
	}

	private void setAuthTokenTimeout(long authTokenTimeout) {
		this.authTokenTimeout = authTokenTimeout;
	}

	/**
	 * get http client
     * @return default httpclient
	 */
	protected HttpClient getHttpClient() {
		if (httpClient == null) {
		    httpClient = new DefaultHttpClient();
		    
			if (super.workUnitState.contains(ConfigurationKeys.SOURCE_USE_PROXY_URL) &&
				! super.workUnitState.getProp(ConfigurationKeys.SOURCE_USE_PROXY_URL).isEmpty()){
				HttpHost proxy = new HttpHost(super.workUnitState.getProp(ConfigurationKeys.SOURCE_USE_PROXY_URL),
					super.workUnitState.getPropAsInt(ConfigurationKeys.SOURCE_USE_PROXY_PORT));
				httpClient.getParams().setParameter(ConnRoutePNames.DEFAULT_PROXY, proxy);
			}
		}
		return httpClient;
	}
	
	@Override
	public void extractMetadata(String schema, String entity, WorkUnit workUnit) throws SchemaException {
		this.log.info("Extract Metadata using Rest Api");
		JsonArray columnArray = new JsonArray();
		String inputQuery = workUnit.getProp(ConfigurationKeys.SOURCE_QUERY);
		List<String> columnListInQuery = null;
		JsonArray array = null;
		if (!Strings.isNullOrEmpty(inputQuery)) {
			columnListInQuery = Utils.getColumnListFromQuery(inputQuery);
		}

		try {
			boolean success = this.getConnection();
			if (!success) {
				throw new SchemaException("Failed to connect.");
			} else {
				this.log.debug("Connected successfully.");
				List<Command> cmds = this.getSchemaMetadata(schema, entity);
				CommandOutput<?, ?> response = this.getResponse(cmds);
				array = this.getSchema(response);

				for (JsonElement columnElement : array) {
					Schema obj = gson.fromJson(columnElement, Schema.class);
					String columnName = obj.getColumnName();

					obj.setWaterMark(this.isWatermarkColumn(workUnit.getProp("extract.delta.fields"), columnName));
					
					if(this.isWatermarkColumn(workUnit.getProp("extract.delta.fields"), columnName)) {
						obj.setNullable(false);
					} else if (this.getPrimarykeyIndex(workUnit.getProp("extract.primary.key.fields"), columnName) == 0){
						// set all columns as nullable except primary key and watermark columns
						obj.setNullable(true);
					}
					
					obj.setPrimaryKey(this.getPrimarykeyIndex(workUnit.getProp("extract.primary.key.fields"), columnName));

					String jsonStr = gson.toJson(obj);
					JsonObject jsonObject = gson.fromJson(jsonStr, JsonObject.class).getAsJsonObject();
					if (inputQuery == null || columnListInQuery == null || (columnListInQuery.size() == 1 && columnListInQuery.get(0).equals("*"))
							|| (columnListInQuery.size() >= 1 && this.isMetadataColumn(columnName, columnListInQuery))) {
						this.columnList.add(columnName);
						columnArray.add(jsonObject);
					}
				}

//				if (inputQuery == null && this.columnList.size() != 0) {
//					this.log.debug("New query with the required column list");
					this.updatedQuery = "SELECT " + Joiner.on(",").join(columnList) + " FROM " + entity;
					log.info(this.updatedQuery);
//				} else {
//					this.log.debug("Query is same as input query");
//					this.updatedQuery = inputQuery;
//				}
				this.log.debug("Schema:" + columnArray);
				this.setOutputSchema(columnArray);
			}

		} catch (Exception e) {
			throw new SchemaException("Failed to get schema using rest api; error - " + e.getMessage(), e);<commit_msg>Fixing restapi<commit_after>package com.linkedin.uif.source.extractor.extract.restapi;

import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.http.HttpEntity;
import org.apache.http.HttpHost;
import org.apache.http.HttpResponse;
import org.apache.http.StatusLine;
import org.apache.http.client.HttpClient;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpRequestBase;
import org.apache.http.conn.params.ConnRoutePNames;
import org.apache.http.impl.client.DefaultHttpClient;
import org.apache.http.util.EntityUtils;

import com.google.common.base.Joiner;
import com.google.common.base.Strings;
import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;

import com.linkedin.uif.configuration.ConfigurationKeys;
import com.linkedin.uif.configuration.WorkUnitState;
import com.linkedin.uif.source.extractor.watermark.Predicate;
import com.linkedin.uif.source.extractor.DataRecordException;
import com.linkedin.uif.source.extractor.exception.HighWatermarkException;
import com.linkedin.uif.source.extractor.exception.RecordCountException;
import com.linkedin.uif.source.extractor.exception.RestApiConnectionException;
import com.linkedin.uif.source.extractor.exception.RestApiProcessingException;
import com.linkedin.uif.source.extractor.exception.SchemaException;
import com.linkedin.uif.source.extractor.extract.BaseExtractor;
import com.linkedin.uif.source.extractor.extract.Command;
import com.linkedin.uif.source.extractor.extract.CommandOutput;
import com.linkedin.uif.source.extractor.extract.SourceSpecificLayer;
import com.linkedin.uif.source.extractor.schema.Schema;
import com.linkedin.uif.source.extractor.utils.Utils;
import com.linkedin.uif.source.workunit.WorkUnit;

/**
 * An implementation of rest api extractor for the sources that are using rest api
 *
 * @param <D> type of data record
 * @param <S> type of schema
 */
public abstract class RestApiExtractor extends BaseExtractor<JsonArray, JsonElement> implements SourceSpecificLayer<JsonArray, JsonElement>, RestApiSpecificLayer {
	private static final Gson gson = new Gson();
	private HttpClient httpClient = null;
	private boolean autoEstablishAuthToken = false;
	private long authTokenTimeout;
	private String accessToken;
	private long createdAt;
	protected String instanceUrl;
	protected String updatedQuery;
	protected Logger log = LoggerFactory.getLogger(RestApiExtractor.class);

	RestApiExtractor(WorkUnitState state) {
		super(state);
	}

	private void setAuthTokenTimeout(long authTokenTimeout) {
		this.authTokenTimeout = authTokenTimeout;
	}

	/**
	 * get http client
     * @return default httpclient
	 */
	protected HttpClient getHttpClient() {
		if (httpClient == null) {
		    httpClient = new DefaultHttpClient();
		    
			if (super.workUnitState.contains(ConfigurationKeys.SOURCE_USE_PROXY_URL) &&
				! super.workUnitState.getProp(ConfigurationKeys.SOURCE_USE_PROXY_URL).isEmpty()){
				HttpHost proxy = new HttpHost(super.workUnitState.getProp(ConfigurationKeys.SOURCE_USE_PROXY_URL),
					super.workUnitState.getPropAsInt(ConfigurationKeys.SOURCE_USE_PROXY_PORT));
				httpClient.getParams().setParameter(ConnRoutePNames.DEFAULT_PROXY, proxy);
			}
		}
		return httpClient;
	}
	
	@Override
	public void extractMetadata(String schema, String entity, WorkUnit workUnit) throws SchemaException {
		this.log.info("Extract Metadata using Rest Api");
		JsonArray columnArray = new JsonArray();
		String inputQuery = workUnit.getProp(ConfigurationKeys.SOURCE_QUERY);
		List<String> columnListInQuery = null;
		JsonArray array = null;
		if (!Strings.isNullOrEmpty(inputQuery)) {
			columnListInQuery = Utils.getColumnListFromQuery(inputQuery);
		}

		try {
			boolean success = this.getConnection();
			if (!success) {
				throw new SchemaException("Failed to connect.");
			} else {
				this.log.debug("Connected successfully.");
				List<Command> cmds = this.getSchemaMetadata(schema, entity);
				CommandOutput<?, ?> response = this.getResponse(cmds);
				array = this.getSchema(response);

				for (JsonElement columnElement : array) {
					Schema obj = gson.fromJson(columnElement, Schema.class);
					String columnName = obj.getColumnName();

					obj.setWaterMark(this.isWatermarkColumn(workUnit.getProp("extract.delta.fields"), columnName));
					
					if(this.isWatermarkColumn(workUnit.getProp("extract.delta.fields"), columnName)) {
						obj.setNullable(false);
					} else if (this.getPrimarykeyIndex(workUnit.getProp("extract.primary.key.fields"), columnName) == 0){
						// set all columns as nullable except primary key and watermark columns
						obj.setNullable(true);
					}
					
					obj.setPrimaryKey(this.getPrimarykeyIndex(workUnit.getProp("extract.primary.key.fields"), columnName));

					String jsonStr = gson.toJson(obj);
					JsonObject jsonObject = gson.fromJson(jsonStr, JsonObject.class).getAsJsonObject();
					if (inputQuery == null || columnListInQuery == null || (columnListInQuery.size() == 1 && columnListInQuery.get(0).equals("*"))
							|| (columnListInQuery.size() >= 1 && this.isMetadataColumn(columnName, columnListInQuery))) {
						this.columnList.add(columnName);
						columnArray.add(jsonObject);
					}
				}

				if (inputQuery == null && this.columnList.size() != 0) {
					this.log.debug("New query with the required column list");
					this.updatedQuery = "SELECT " + Joiner.on(",").join(columnList) + " FROM " + entity;
					log.info(this.updatedQuery);
				} else {
					this.log.debug("Query is same as input query");
					this.updatedQuery = inputQuery;
				}
				this.log.debug("Schema:" + columnArray);
				this.setOutputSchema(columnArray);
			}

		} catch (Exception e) {
			throw new SchemaException("Failed to get schema using rest api; error - " + e.getMessage(), e);
{'fraction_non_alphanumeric': 0.11363061620163534, 'fraction_numerical': 0.0011801399308775183, 'mean_word_length': 4.908366533864542, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4268932', 'n_tokens_mistral': 3940, 'n_tokens_neox': 3558, 'n_words': 678}